"""
AtomizerSpec v2.0 Pydantic Models

These models match the JSON Schema at optimization_engine/schemas/atomizer_spec_v2.json
They provide validation and type safety for the unified configuration system.
"""

from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, Field, field_validator, model_validator
import re


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

class SpecCreatedBy(str, Enum):
    """Who/what created the spec."""
    CANVAS = "canvas"
    CLAUDE = "claude"
    API = "api"
    MIGRATION = "migration"
    MANUAL = "manual"


class SolverType(str, Enum):
    """Supported solver types."""
    NASTRAN = "nastran"
    NX_NASTRAN = "NX_Nastran"
    ABAQUS = "abaqus"


class SubcaseType(str, Enum):
    """Subcase analysis types."""
    STATIC = "static"
    MODAL = "modal"
    THERMAL = "thermal"
    BUCKLING = "buckling"


class DesignVariableType(str, Enum):
    """Design variable types."""
    CONTINUOUS = "continuous"
    INTEGER = "integer"
    CATEGORICAL = "categorical"


class ExtractorType(str, Enum):
    """Physics extractor types."""
    DISPLACEMENT = "displacement"
    FREQUENCY = "frequency"
    STRESS = "stress"
    MASS = "mass"
    MASS_EXPRESSION = "mass_expression"
    ZERNIKE_OPD = "zernike_opd"
    ZERNIKE_CSV = "zernike_csv"
    TEMPERATURE = "temperature"
    CUSTOM_FUNCTION = "custom_function"


class OptimizationDirection(str, Enum):
    """Optimization direction."""
    MINIMIZE = "minimize"
    MAXIMIZE = "maximize"


class ConstraintType(str, Enum):
    """Constraint types."""
    HARD = "hard"
    SOFT = "soft"


class ConstraintOperator(str, Enum):
    """Constraint comparison operators."""
    LE = "<="
    GE = ">="
    LT = "<"
    GT = ">"
    EQ = "=="


class PenaltyMethod(str, Enum):
    """Penalty methods for constraints."""
    LINEAR = "linear"
    QUADRATIC = "quadratic"
    EXPONENTIAL = "exponential"


class AlgorithmType(str, Enum):
    """Optimization algorithm types."""
    TPE = "TPE"
    CMA_ES = "CMA-ES"
    NSGA_II = "NSGA-II"
    RANDOM_SEARCH = "RandomSearch"
    SAT_V3 = "SAT_v3"
    GP_BO = "GP-BO"


class SurrogateType(str, Enum):
    """Surrogate model types."""
    MLP = "MLP"
    GNN = "GNN"
    ENSEMBLE = "ensemble"


# ============================================================================
# Position Model
# ============================================================================

class CanvasPosition(BaseModel):
    """Canvas position for nodes."""
    x: float = 0
    y: float = 0


# ============================================================================
# Meta Models
# ============================================================================

class SpecMeta(BaseModel):
    """Metadata about the spec."""
    version: str = Field(
        ...,
        pattern=r"^2\.\d+$",
        description="Schema version (e.g., '2.0')"
    )
    created: Optional[datetime] = Field(
        default=None,
        description="When the spec was created"
    )
    modified: Optional[datetime] = Field(
        default=None,
        description="When the spec was last modified"
    )
    created_by: Optional[SpecCreatedBy] = Field(
        default=None,
        description="Who/what created the spec"
    )
    modified_by: Optional[str] = Field(
        default=None,
        description="Who/what last modified the spec"
    )
    study_name: str = Field(
        ...,
        min_length=3,
        max_length=100,
        pattern=r"^[a-z0-9_]+$",
        description="Unique study identifier (snake_case)"
    )
    description: Optional[str] = Field(
        default=None,
        max_length=1000,
        description="Human-readable description"
    )
    tags: Optional[List[str]] = Field(
        default=None,
        description="Tags for categorization"
    )
    engineering_context: Optional[str] = Field(
        default=None,
        description="Real-world engineering context"
    )


# ============================================================================
# Model Configuration Models
# ============================================================================

class NxPartConfig(BaseModel):
    """NX geometry part file configuration."""
    path: Optional[str] = Field(default=None, description="Path to .prt file")
    hash: Optional[str] = Field(default=None, description="File hash for change detection")
    idealized_part: Optional[str] = Field(default=None, description="Idealized part filename (_i.prt)")


class FemConfig(BaseModel):
    """FEM mesh file configuration."""
    path: Optional[str] = Field(default=None, description="Path to .fem file")
    element_count: Optional[int] = Field(default=None, description="Number of elements")
    node_count: Optional[int] = Field(default=None, description="Number of nodes")


class Subcase(BaseModel):
    """Simulation subcase definition."""
    id: int
    name: Optional[str] = None
    type: Optional[SubcaseType] = None


class SimConfig(BaseModel):
    """Simulation file configuration."""
    path: str = Field(..., description="Path to .sim file")
    solver: SolverType = Field(..., description="Solver type")
    solution_type: Optional[str] = Field(
        default=None,
        pattern=r"^SOL\d+$",
        description="Solution type (e.g., SOL101)"
    )
    subcases: Optional[List[Subcase]] = Field(default=None, description="Defined subcases")


class NxSettings(BaseModel):
    """NX runtime settings."""
    nx_install_path: Optional[str] = None
    simulation_timeout_s: Optional[int] = Field(default=None, ge=60, le=7200)
    auto_start_nx: Optional[bool] = None


class ModelConfig(BaseModel):
    """NX model files and configuration."""
    nx_part: Optional[NxPartConfig] = None
    fem: Optional[FemConfig] = None
    sim: SimConfig
    nx_settings: Optional[NxSettings] = None


# ============================================================================
# Design Variable Models
# ============================================================================

class DesignVariableBounds(BaseModel):
    """Design variable bounds."""
    min: float
    max: float

    @model_validator(mode='after')
    def validate_bounds(self) -> 'DesignVariableBounds':
        if self.min >= self.max:
            raise ValueError(f"min ({self.min}) must be less than max ({self.max})")
        return self


class DesignVariable(BaseModel):
    """A design variable to optimize."""
    id: str = Field(
        ...,
        pattern=r"^dv_\d{3}$",
        description="Unique identifier (pattern: dv_XXX)"
    )
    name: str = Field(..., description="Human-readable name")
    expression_name: str = Field(
        ...,
        pattern=r"^[a-zA-Z_][a-zA-Z0-9_]*$",
        description="NX expression name (must match model)"
    )
    type: DesignVariableType = Field(..., description="Variable type")
    bounds: DesignVariableBounds = Field(..., description="Value bounds")
    baseline: Optional[float] = Field(default=None, description="Current/initial value")
    units: Optional[str] = Field(default=None, description="Physical units (mm, deg, etc.)")
    step: Optional[float] = Field(default=None, description="Step size for integer/discrete")
    enabled: bool = Field(default=True, description="Whether to include in optimization")
    description: Optional[str] = None
    canvas_position: Optional[CanvasPosition] = None


# ============================================================================
# Extractor Models
# ============================================================================

class ExtractorConfig(BaseModel):
    """Type-specific extractor configuration."""
    inner_radius_mm: Optional[float] = None
    outer_radius_mm: Optional[float] = None
    n_modes: Optional[int] = None
    filter_low_orders: Optional[int] = None
    displacement_unit: Optional[str] = None
    reference_subcase: Optional[int] = None
    expression_name: Optional[str] = None
    mode_number: Optional[int] = None
    element_type: Optional[str] = None
    result_type: Optional[str] = None
    metric: Optional[str] = None

    class Config:
        extra = "allow"  # Allow additional fields for flexibility


class CustomFunction(BaseModel):
    """Custom function definition for custom_function extractors."""
    name: Optional[str] = Field(default=None, description="Function name")
    module: Optional[str] = Field(default=None, description="Python module path")
    signature: Optional[str] = Field(default=None, description="Function signature")
    source_code: Optional[str] = Field(default=None, description="Python source code")


class ExtractorOutput(BaseModel):
    """Output definition for an extractor."""
    name: str = Field(..., description="Output name (used by objectives/constraints)")
    metric: Optional[str] = Field(default=None, description="Specific metric (max, total, rms, etc.)")
    subcase: Optional[int] = Field(default=None, description="Subcase ID for this output")
    units: Optional[str] = None


class Extractor(BaseModel):
    """Physics extractor that computes outputs from FEA."""
    id: str = Field(
        ...,
        pattern=r"^ext_\d{3}$",
        description="Unique identifier (pattern: ext_XXX)"
    )
    name: str = Field(..., description="Human-readable name")
    type: ExtractorType = Field(..., description="Extractor type")
    builtin: bool = Field(default=True, description="Whether this is a built-in extractor")
    config: Optional[ExtractorConfig] = Field(default=None, description="Type-specific configuration")
    function: Optional[CustomFunction] = Field(
        default=None,
        description="Custom function definition (for custom_function type)"
    )
    outputs: List[ExtractorOutput] = Field(..., min_length=1, description="Output values")
    canvas_position: Optional[CanvasPosition] = None

    @model_validator(mode='after')
    def validate_custom_function(self) -> 'Extractor':
        if self.type == ExtractorType.CUSTOM_FUNCTION and self.function is None:
            raise ValueError("custom_function extractor requires function definition")
        return self


# ============================================================================
# Objective Models
# ============================================================================

class ObjectiveSource(BaseModel):
    """Source reference for objective value."""
    extractor_id: str = Field(..., description="Reference to extractor")
    output_name: str = Field(..., description="Which output from the extractor")


class Objective(BaseModel):
    """Optimization objective."""
    id: str = Field(
        ...,
        pattern=r"^obj_\d{3}$",
        description="Unique identifier (pattern: obj_XXX)"
    )
    name: str = Field(..., description="Human-readable name")
    direction: OptimizationDirection = Field(..., description="Optimization direction")
    weight: float = Field(default=1.0, ge=0, description="Weight for weighted sum")
    source: ObjectiveSource = Field(..., description="Where the value comes from")
    target: Optional[float] = Field(default=None, description="Target value (for goal programming)")
    units: Optional[str] = None
    description: Optional[str] = None
    canvas_position: Optional[CanvasPosition] = None


# ============================================================================
# Constraint Models
# ============================================================================

class ConstraintSource(BaseModel):
    """Source reference for constraint value."""
    extractor_id: str
    output_name: str


class PenaltyConfig(BaseModel):
    """Penalty method configuration for constraints."""
    method: Optional[PenaltyMethod] = None
    weight: Optional[float] = None
    margin: Optional[float] = Field(default=None, description="Soft margin before penalty kicks in")


class Constraint(BaseModel):
    """Hard or soft constraint."""
    id: str = Field(
        ...,
        pattern=r"^con_\d{3}$",
        description="Unique identifier (pattern: con_XXX)"
    )
    name: str
    type: ConstraintType = Field(..., description="Constraint type")
    operator: ConstraintOperator = Field(..., description="Comparison operator")
    threshold: float = Field(..., description="Constraint threshold value")
    source: ConstraintSource = Field(..., description="Where the value comes from")
    penalty_config: Optional[PenaltyConfig] = None
    description: Optional[str] = None
    canvas_position: Optional[CanvasPosition] = None


# ============================================================================
# Optimization Models
# ============================================================================

class AlgorithmConfig(BaseModel):
    """Algorithm-specific settings."""
    population_size: Optional[int] = None
    n_generations: Optional[int] = None
    mutation_prob: Optional[float] = None
    crossover_prob: Optional[float] = None
    seed: Optional[int] = None
    n_startup_trials: Optional[int] = None
    sigma0: Optional[float] = None

    class Config:
        extra = "allow"  # Allow additional algorithm-specific fields


class Algorithm(BaseModel):
    """Optimization algorithm configuration."""
    type: AlgorithmType
    config: Optional[AlgorithmConfig] = None


class OptimizationBudget(BaseModel):
    """Computational budget for optimization."""
    max_trials: Optional[int] = Field(default=None, ge=1, le=10000)
    max_time_hours: Optional[float] = None
    convergence_patience: Optional[int] = Field(
        default=None,
        description="Stop if no improvement for N trials"
    )


class SurrogateConfig(BaseModel):
    """Neural surrogate model configuration."""
    n_models: Optional[int] = None
    architecture: Optional[List[int]] = None
    train_every_n_trials: Optional[int] = None
    min_training_samples: Optional[int] = None
    acquisition_candidates: Optional[int] = None
    fea_validations_per_round: Optional[int] = None


class Surrogate(BaseModel):
    """Surrogate model settings."""
    enabled: Optional[bool] = None
    type: Optional[SurrogateType] = None
    config: Optional[SurrogateConfig] = None


class OptimizationConfig(BaseModel):
    """Optimization algorithm configuration."""
    algorithm: Algorithm
    budget: OptimizationBudget
    surrogate: Optional[Surrogate] = None
    canvas_position: Optional[CanvasPosition] = None


# ============================================================================
# Workflow Models
# ============================================================================

class WorkflowStage(BaseModel):
    """A stage in a multi-stage optimization workflow."""
    id: str
    name: str
    algorithm: Optional[str] = None
    trials: Optional[int] = None
    purpose: Optional[str] = None


class WorkflowTransition(BaseModel):
    """Transition between workflow stages."""
    from_: str = Field(..., alias="from")
    to: str
    condition: Optional[str] = None

    class Config:
        populate_by_name = True


class Workflow(BaseModel):
    """Multi-stage optimization workflow."""
    stages: Optional[List[WorkflowStage]] = None
    transitions: Optional[List[WorkflowTransition]] = None


# ============================================================================
# Reporting Models
# ============================================================================

class InsightConfig(BaseModel):
    """Insight-specific configuration."""
    include_html: Optional[bool] = None
    show_pareto_evolution: Optional[bool] = None

    class Config:
        extra = "allow"


class Insight(BaseModel):
    """Reporting insight definition."""
    type: Optional[str] = None
    for_trials: Optional[str] = None
    config: Optional[InsightConfig] = None


class ReportingConfig(BaseModel):
    """Reporting configuration."""
    auto_report: Optional[bool] = None
    report_triggers: Optional[List[str]] = None
    insights: Optional[List[Insight]] = None


# ============================================================================
# Canvas Models
# ============================================================================

class CanvasViewport(BaseModel):
    """Canvas viewport settings."""
    x: float = 0
    y: float = 0
    zoom: float = 1.0


class CanvasEdge(BaseModel):
    """Connection between canvas nodes."""
    source: str
    target: str
    sourceHandle: Optional[str] = None
    targetHandle: Optional[str] = None


class CanvasGroup(BaseModel):
    """Grouping of canvas nodes."""
    id: str
    name: str
    node_ids: List[str]


class CanvasConfig(BaseModel):
    """Canvas UI state (persisted for reconstruction)."""
    layout_version: Optional[str] = None
    viewport: Optional[CanvasViewport] = None
    edges: Optional[List[CanvasEdge]] = None
    groups: Optional[List[CanvasGroup]] = None


# ============================================================================
# Main AtomizerSpec Model
# ============================================================================

class AtomizerSpec(BaseModel):
    """
    AtomizerSpec v2.0 - The unified configuration schema for Atomizer optimization studies.

    This is the single source of truth used by:
    - Canvas UI (rendering and editing)
    - Backend API (validation and storage)
    - Claude Assistant (reading and modifying)
    - Optimization Engine (execution)
    """
    meta: SpecMeta = Field(..., description="Metadata about the spec")
    model: ModelConfig = Field(..., description="NX model files and configuration")
    design_variables: List[DesignVariable] = Field(
        ...,
        min_length=1,
        max_length=50,
        description="Design variables to optimize"
    )
    extractors: List[Extractor] = Field(
        ...,
        min_length=1,
        description="Physics extractors"
    )
    objectives: List[Objective] = Field(
        ...,
        min_length=1,
        max_length=5,
        description="Optimization objectives"
    )
    constraints: Optional[List[Constraint]] = Field(
        default=None,
        description="Hard and soft constraints"
    )
    optimization: OptimizationConfig = Field(..., description="Algorithm configuration")
    workflow: Optional[Workflow] = Field(default=None, description="Multi-stage workflow")
    reporting: Optional[ReportingConfig] = Field(default=None, description="Reporting config")
    canvas: Optional[CanvasConfig] = Field(default=None, description="Canvas UI state")

    @model_validator(mode='after')
    def validate_references(self) -> 'AtomizerSpec':
        """Validate that all references are valid."""
        # Collect valid extractor IDs and their outputs
        extractor_outputs: Dict[str, set] = {}
        for ext in self.extractors:
            extractor_outputs[ext.id] = {o.name for o in ext.outputs}

        # Validate objective sources
        for obj in self.objectives:
            if obj.source.extractor_id not in extractor_outputs:
                raise ValueError(
                    f"Objective '{obj.name}' references unknown extractor: {obj.source.extractor_id}"
                )
            if obj.source.output_name not in extractor_outputs[obj.source.extractor_id]:
                raise ValueError(
                    f"Objective '{obj.name}' references unknown output: {obj.source.output_name}"
                )

        # Validate constraint sources
        if self.constraints:
            for con in self.constraints:
                if con.source.extractor_id not in extractor_outputs:
                    raise ValueError(
                        f"Constraint '{con.name}' references unknown extractor: {con.source.extractor_id}"
                    )
                if con.source.output_name not in extractor_outputs[con.source.extractor_id]:
                    raise ValueError(
                        f"Constraint '{con.name}' references unknown output: {con.source.output_name}"
                    )

        return self

    def get_enabled_design_variables(self) -> List[DesignVariable]:
        """Return only enabled design variables."""
        return [dv for dv in self.design_variables if dv.enabled]

    def get_extractor_by_id(self, extractor_id: str) -> Optional[Extractor]:
        """Find an extractor by ID."""
        for ext in self.extractors:
            if ext.id == extractor_id:
                return ext
        return None

    def get_objective_by_id(self, objective_id: str) -> Optional[Objective]:
        """Find an objective by ID."""
        for obj in self.objectives:
            if obj.id == objective_id:
                return obj
        return None

    def get_constraint_by_id(self, constraint_id: str) -> Optional[Constraint]:
        """Find a constraint by ID."""
        if not self.constraints:
            return None
        for con in self.constraints:
            if con.id == constraint_id:
                return con
        return None

    def has_custom_extractors(self) -> bool:
        """Check if spec has any custom function extractors."""
        return any(ext.type == ExtractorType.CUSTOM_FUNCTION for ext in self.extractors)

    def is_multi_objective(self) -> bool:
        """Check if this is a multi-objective optimization."""
        return len(self.objectives) > 1


# ============================================================================
# Validation Response Models
# ============================================================================

class ValidationError(BaseModel):
    """A validation error."""
    type: str  # 'schema', 'semantic', 'reference'
    path: List[str]
    message: str


class ValidationWarning(BaseModel):
    """A validation warning."""
    type: str
    path: List[str]
    message: str


class ValidationSummary(BaseModel):
    """Summary of spec contents."""
    design_variables: int
    extractors: int
    objectives: int
    constraints: int
    custom_functions: int


class ValidationReport(BaseModel):
    """Full validation report."""
    valid: bool
    errors: List[ValidationError]
    warnings: List[ValidationWarning]
    summary: ValidationSummary