"""Simulation Plan Schema V2 (P3-M1).

Extends V1 with multi-fidelity strategy, search strategy,
calibration policy, and acceptance criteria per third-party review.

Backward compatible: V1 plans without these fields use defaults.
"""
from enum import Enum
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field


# ============================================================
# Enums
# ============================================================

class StrategyMode(str, Enum):
    """Simulation strategy mode."""
    FAST_FEASIBLE = "fast_feasible"       # Default: constrained Bayesian / active learning
    PARETO_EXPLORATION = "pareto_exploration"  # Morris + LHS + Kriging + NSGA-II
    HIGH_FIDELITY_VALIDATION = "high_fidelity_validation"  # L2 + L3 verification
    ROBUSTNESS_CHECK = "robustness_check"  # Tolerance / disturbance analysis


class FidelityLevel(str, Enum):
    """Multi-fidelity levels L0-L4."""
    L0_ANALYTIC = "L0_analytic"                    # Analytic formulas + rule engine
    L1_MOTORCAD_EMAG = "L1_motorcad_emag"         # Motor-CAD fast EM model
    L2_MOTORCAD_LAB_THERM = "L2_motorcad_lab_therm"  # Motor-CAD Lab/Therm/Mech
    L3_MAXWELL_3D = "L3_maxwell_3d"               # Maxwell 3D / JMAG high-fidelity
    L4_ROBUSTNESS = "L4_robustness"                # Tolerance / prototype data


class SearchMethod(str, Enum):
    """Search/optimization method."""
    CONSTRAINED_BAYESIAN = "constrained_bayesian"  # Default: feasibility-first
    ACTIVE_LEARNING = "active_learning"             # Uncertainty-based sampling
    LHS_KRIGING_NSGA2 = "lhs_kriging_nsga2"        # Global Pareto exploration
    GRID_SCAN = "grid_scan"                          # Fixed full factorial
    TRUST_REGION = "trust_region"                    # Local trust region search


class ConfidenceGrade(str, Enum):
    """Result confidence grade A/B/C/D."""
    A = "A"  # Full multi-physics + high-fidelity + robustness, design freeze ready
    B = "B"  # Motor-CAD + at least one high-fidelity check, candidate ready
    C = "C"  # Motor-CAD only, internal discussion only, no external commitment
    D = "D"  # Analytic / surrogate only, reference only


class ConvergenceStatus(str, Enum):
    """Six types of convergence status (per review)."""
    # Solver convergence
    SOLVER_PASS = "SOLVER_PASS"
    SOLVER_FAIL = "SOLVER_FAIL"
    # Hard constraint convergence
    FEASIBLE = "FEASIBLE"
    INFEASIBLE = "INFEASIBLE"
    # Optimization convergence
    CONVERGED = "CONVERGED"
    STALLED = "STALLED"
    # Surrogate model trust
    MODEL_TRUSTED = "MODEL_TRUSTED"
    MODEL_UNCERTAIN = "MODEL_UNCERTAIN"
    # Cross-tool consistency
    HF_PASS = "HF_PASS"
    HF_FAIL = "HF_FAIL"
    # Robustness
    ROBUST = "ROBUST"
    FRAGILE = "FRAGILE"


# ============================================================
# Strategy Sub-models
# ============================================================

class FidelityStrategy(BaseModel):
    """Multi-fidelity execution strategy."""
    levels: List[FidelityLevel] = Field(
        default_factory=lambda: [
            FidelityLevel.L0_ANALYTIC,
            FidelityLevel.L1_MOTORCAD_EMAG,
        ],
        description="Enabled fidelity levels in execution order"
    )
    upgrade_rule: str = Field(
        default="top_candidates_only",
        description="Rule for upgrading to higher fidelity: top_candidates_only / threshold_based / all"
    )
    max_candidates_for_l3: int = Field(
        default=3, ge=1, le=10,
        description="Max candidates to send to L3 (Maxwell/JMAG)"
    )
    max_solver_cost_per_level: Optional[Dict[str, int]] = Field(
        default=None,
        description="Max solver calls per fidelity level, e.g. {'L1': 80, 'L2': 10}"
    )


class SearchStrategy(BaseModel):
    """Adaptive search strategy."""
    method: SearchMethod = Field(
        default=SearchMethod.CONSTRAINED_BAYESIAN,
        description="Search/optimization method"
    )
    initial_samples: int = Field(
        default=16, ge=4, le=100,
        description="Number of initial samples (LHS or from experience)"
    )
    batch_size: int = Field(
        default=4, ge=1, le=16,
        description="Number of points per adaptive batch"
    )
    max_solver_calls: int = Field(
        default=80, ge=10, le=500,
        description="Maximum total solver calls"
    )
    local_trust_region: bool = Field(
        default=True,
        description="Enable local trust region refinement after feasible region found"
    )
    trust_region_radius: Optional[float] = Field(
        default=None,
        description="Initial trust region radius as fraction of parameter range"
    )
    use_experience_seeds: bool = Field(
        default=True,
        description="Use similar experience cases as initial seeds"
    )


class CalibrationPolicy(BaseModel):
    """Cross-tool calibration policy (Motor-CAD vs Maxwell/JMAG)."""
    enabled: bool = Field(default=False, description="Enable cross-tool calibration")
    cross_tool_metrics: List[str] = Field(
        default_factory=lambda: ["torque_nm", "efficiency_pct", "axial_force_n"],
        description="Metrics to compare across tools"
    )
    tolerance: Dict[str, float] = Field(
        default_factory=lambda: {"torque_pct": 5.0, "efficiency_point": 1.0, "axial_force_pct": 10.0},
        description="Acceptable tolerance for cross-tool deviation"
    )
    correction_method: str = Field(
        default="additive",
        description="Correction method: additive / multiplicative / co-kriging"
    )
    feedback_to_surrogate: bool = Field(
        default=True,
        description="Feed calibration coefficients back to surrogate model and objective"
    )


class AcceptanceCriteria(BaseModel):
    """Acceptance criteria for convergence and validation."""
    hard_constraints: List[str] = Field(
        default_factory=list,
        description="Hard constraint expressions, e.g. ['torque_nm >= 10', 'temperature_c <= 120']"
    )
    soft_objectives: Optional[List[str]] = Field(
        default=None,
        description="Soft objective expressions for optimization"
    )
    cross_tool_tolerance: Optional[Dict[str, float]] = Field(
        default=None,
        description="Override cross-tool tolerance"
    )
    surrogate_max_uncertainty: float = Field(
        default=0.05, ge=0.01, le=0.5,
        description="Max surrogate model uncertainty for MODEL_TRUSTED status"
    )
    robustness_required: bool = Field(
        default=False,
        description="Require robustness check before final acceptance"
    )
    min_confidence_grade: ConfidenceGrade = Field(
        default=ConfidenceGrade.C,
        description="Minimum confidence grade for plan acceptance"
    )


class ParallelExecution(BaseModel):
    """Parallel execution configuration."""
    max_instances: int = Field(
        default=1, ge=1, le=8,
        description="Max parallel Motor-CAD instances"
    )
    model_copy_strategy: str = Field(
        default="per_instance",
        description="Model file copy strategy: per_instance / shared_readonly"
    )
    license_fail_policy: str = Field(
        default="queue_retry",
        description="Policy on license failure: queue_retry / fail_fast / reduce_instances"
    )


# ============================================================
# Schema V2 Main Model
# ============================================================

class SimulationPlanSchemaV2(BaseModel):
    """Extended simulation plan schema (V2) per third-party review.

    All new fields are optional for backward compatibility with V1 plans.
    """
    schema_version: str = Field(default="2.0", description="Schema version")
    strategy_mode: StrategyMode = Field(
        default=StrategyMode.FAST_FEASIBLE,
        description="Simulation strategy mode"
    )
    fidelity_strategy: Optional[FidelityStrategy] = Field(
        default=None,
        description="Multi-fidelity execution strategy"
    )
    search_strategy: Optional[SearchStrategy] = Field(
        default=None,
        description="Adaptive search strategy"
    )
    calibration_policy: Optional[CalibrationPolicy] = Field(
        default=None,
        description="Cross-tool calibration policy"
    )
    acceptance_criteria: Optional[AcceptanceCriteria] = Field(
        default=None,
        description="Acceptance criteria"
    )
    parallel_execution: Optional[ParallelExecution] = Field(
        default=None,
        description="Parallel execution configuration"
    )

    # V1 fields (preserved)
    model_path: Optional[str] = None
    topology: Optional[str] = None
    variables: List[Dict[str, Any]] = Field(default_factory=list)
    cases: List[Dict[str, Any]] = Field(default_factory=list)
    boundary_conditions: Optional[Dict[str, Any]] = None

    def get_effective_fidelity(self) -> FidelityStrategy:
        """Get fidelity strategy with defaults applied."""
        return self.fidelity_strategy or FidelityStrategy()

    def get_effective_search(self) -> SearchStrategy:
        """Get search strategy with defaults applied."""
        return self.search_strategy or SearchStrategy()

    def get_effective_acceptance(self) -> AcceptanceCriteria:
        """Get acceptance criteria with defaults applied."""
        return self.acceptance_criteria or AcceptanceCriteria()