"""
Study Blueprint

Data structures for the study blueprint - the validated configuration
ready for study generation.
"""

from dataclasses import dataclass, field, asdict
from typing import Dict, List, Any, Optional
import json


@dataclass
class DesignVariable:
    """Design variable specification."""
    parameter: str
    current_value: float
    min_value: float
    max_value: float
    units: Optional[str] = None
    is_integer: bool = False

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def to_config_format(self) -> Dict[str, Any]:
        """Convert to optimization_config.json format."""
        return {
            "expression_name": self.parameter,
            "bounds": [self.min_value, self.max_value],
            "units": self.units or "",
            "is_integer": self.is_integer,
        }


@dataclass
class Objective:
    """Optimization objective specification."""
    name: str
    goal: str  # minimize, maximize, target
    extractor: str  # Extractor ID (e.g., "E1", "E4")
    extractor_name: Optional[str] = None
    extractor_params: Optional[Dict[str, Any]] = None
    weight: float = 1.0
    target_value: Optional[float] = None  # For target objectives

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def to_config_format(self) -> Dict[str, Any]:
        """Convert to optimization_config.json format."""
        config = {
            "name": self.name,
            "type": self.goal,
            "extractor": self.extractor,
            "weight": self.weight,
        }
        if self.extractor_params:
            config["extractor_params"] = self.extractor_params
        if self.target_value is not None:
            config["target"] = self.target_value
        return config


@dataclass
class Constraint:
    """Optimization constraint specification."""
    name: str
    constraint_type: str  # max, min
    threshold: float
    extractor: str  # Extractor ID
    extractor_name: Optional[str] = None
    extractor_params: Optional[Dict[str, Any]] = None
    is_hard: bool = True
    penalty_weight: float = 1000.0  # For soft constraints

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)

    def to_config_format(self) -> Dict[str, Any]:
        """Convert to optimization_config.json format."""
        config = {
            "name": self.name,
            "type": self.constraint_type,
            "threshold": self.threshold,
            "extractor": self.extractor,
            "hard": self.is_hard,
        }
        if self.extractor_params:
            config["extractor_params"] = self.extractor_params
        if not self.is_hard:
            config["penalty_weight"] = self.penalty_weight
        return config


@dataclass
class StudyBlueprint:
    """
    Complete study blueprint ready for generation.

    This is the validated configuration that will be used to create
    the study files (optimization_config.json, run_optimization.py, etc.)
    """
    # Study metadata
    study_name: str
    study_description: str = ""
    interview_session_id: str = ""

    # Model paths
    model_path: str = ""
    sim_path: str = ""
    fem_path: str = ""

    # Design space
    design_variables: List[DesignVariable] = field(default_factory=list)

    # Optimization goals
    objectives: List[Objective] = field(default_factory=list)
    constraints: List[Constraint] = field(default_factory=list)

    # Optimization settings
    protocol: str = "protocol_10_single"  # or "protocol_11_multi"
    n_trials: int = 100
    sampler: str = "TPE"
    use_neural_acceleration: bool = False

    # Solver settings
    solver_config: Dict[str, Any] = field(default_factory=dict)
    solve_all_solutions: bool = True

    # Extractors configuration
    extractors_config: Dict[str, Any] = field(default_factory=dict)

    # Validation
    warnings_acknowledged: List[str] = field(default_factory=list)
    baseline_validated: bool = False
    baseline_results: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "study_name": self.study_name,
            "study_description": self.study_description,
            "interview_session_id": self.interview_session_id,
            "model_path": self.model_path,
            "sim_path": self.sim_path,
            "fem_path": self.fem_path,
            "design_variables": [dv.to_dict() for dv in self.design_variables],
            "objectives": [obj.to_dict() for obj in self.objectives],
            "constraints": [con.to_dict() for con in self.constraints],
            "protocol": self.protocol,
            "n_trials": self.n_trials,
            "sampler": self.sampler,
            "use_neural_acceleration": self.use_neural_acceleration,
            "solver_config": self.solver_config,
            "solve_all_solutions": self.solve_all_solutions,
            "extractors_config": self.extractors_config,
            "warnings_acknowledged": self.warnings_acknowledged,
            "baseline_validated": self.baseline_validated,
            "baseline_results": self.baseline_results,
        }

    def to_json(self) -> str:
        """Serialize to JSON string."""
        return json.dumps(self.to_dict(), indent=2)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "StudyBlueprint":
        """Create from dictionary."""
        design_variables = [
            DesignVariable(**dv) for dv in data.get("design_variables", [])
        ]
        objectives = [
            Objective(**obj) for obj in data.get("objectives", [])
        ]
        constraints = [
            Constraint(**con) for con in data.get("constraints", [])
        ]

        return cls(
            study_name=data.get("study_name", ""),
            study_description=data.get("study_description", ""),
            interview_session_id=data.get("interview_session_id", ""),
            model_path=data.get("model_path", ""),
            sim_path=data.get("sim_path", ""),
            fem_path=data.get("fem_path", ""),
            design_variables=design_variables,
            objectives=objectives,
            constraints=constraints,
            protocol=data.get("protocol", "protocol_10_single"),
            n_trials=data.get("n_trials", 100),
            sampler=data.get("sampler", "TPE"),
            use_neural_acceleration=data.get("use_neural_acceleration", False),
            solver_config=data.get("solver_config", {}),
            solve_all_solutions=data.get("solve_all_solutions", True),
            extractors_config=data.get("extractors_config", {}),
            warnings_acknowledged=data.get("warnings_acknowledged", []),
            baseline_validated=data.get("baseline_validated", False),
            baseline_results=data.get("baseline_results"),
        )

    def to_config_json(self) -> Dict[str, Any]:
        """
        Convert to optimization_config.json format.

        This is the format expected by the optimization runner.
        """
        config = {
            "study_name": self.study_name,
            "description": self.study_description,
            "version": "2.0",

            "model": {
                "part_file": self.model_path,
                "sim_file": self.sim_path,
                "fem_file": self.fem_path,
            },

            "design_variables": [
                dv.to_config_format() for dv in self.design_variables
            ],

            "objectives": [
                obj.to_config_format() for obj in self.objectives
            ],

            "constraints": [
                con.to_config_format() for con in self.constraints
            ],

            "optimization": {
                "n_trials": self.n_trials,
                "sampler": self.sampler,
                "protocol": self.protocol,
                "neural_acceleration": self.use_neural_acceleration,
            },

            "solver": {
                "solve_all": self.solve_all_solutions,
                **self.solver_config,
            },

            "extractors": self.extractors_config,

            "_metadata": {
                "interview_session_id": self.interview_session_id,
                "warnings_acknowledged": self.warnings_acknowledged,
                "baseline_validated": self.baseline_validated,
            }
        }

        return config

    def to_markdown(self) -> str:
        """Generate human-readable markdown summary."""
        lines = []

        lines.append(f"# Study Blueprint: {self.study_name}")
        lines.append("")

        if self.study_description:
            lines.append(f"**Description**: {self.study_description}")
            lines.append("")

        # Design Variables
        lines.append(f"## Design Variables ({len(self.design_variables)})")
        lines.append("")
        lines.append("| Parameter | Current | Min | Max | Units |")
        lines.append("|-----------|---------|-----|-----|-------|")
        for dv in self.design_variables:
            lines.append(f"| {dv.parameter} | {dv.current_value} | {dv.min_value} | {dv.max_value} | {dv.units or '-'} |")
        lines.append("")

        # Objectives
        lines.append(f"## Objectives ({len(self.objectives)})")
        lines.append("")
        lines.append("| Name | Goal | Extractor | Weight |")
        lines.append("|------|------|-----------|--------|")
        for obj in self.objectives:
            lines.append(f"| {obj.name} | {obj.goal} | {obj.extractor} | {obj.weight} |")
        lines.append("")

        # Constraints
        if self.constraints:
            lines.append(f"## Constraints ({len(self.constraints)})")
            lines.append("")
            lines.append("| Name | Type | Threshold | Extractor | Hard? |")
            lines.append("|------|------|-----------|-----------|-------|")
            for con in self.constraints:
                op = "<=" if con.constraint_type == "max" else ">="
                lines.append(f"| {con.name} | {op} | {con.threshold} | {con.extractor} | {'Yes' if con.is_hard else 'No'} |")
            lines.append("")

        # Settings
        lines.append("## Optimization Settings")
        lines.append("")
        lines.append(f"- **Protocol**: {self.protocol}")
        lines.append(f"- **Trials**: {self.n_trials}")
        lines.append(f"- **Sampler**: {self.sampler}")
        lines.append(f"- **Neural Acceleration**: {'Enabled' if self.use_neural_acceleration else 'Disabled'}")
        lines.append("")

        # Validation
        lines.append("## Validation")
        lines.append("")
        lines.append(f"- **Baseline Validated**: {'Yes' if self.baseline_validated else 'No'}")
        if self.warnings_acknowledged:
            lines.append(f"- **Warnings Acknowledged**: {len(self.warnings_acknowledged)}")
            for w in self.warnings_acknowledged:
                lines.append(f"  - {w}")
        lines.append("")

        return "\n".join(lines)

    def validate(self) -> List[str]:
        """
        Validate blueprint completeness.

        Returns:
            List of validation errors (empty if valid)
        """
        errors = []

        if not self.study_name:
            errors.append("Study name is required")

        if not self.design_variables:
            errors.append("At least one design variable is required")

        if not self.objectives:
            errors.append("At least one objective is required")

        for dv in self.design_variables:
            if dv.min_value >= dv.max_value:
                errors.append(f"Invalid bounds for {dv.parameter}: min >= max")

        return errors

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

    def get_objective_count(self) -> int:
        """Get number of objectives."""
        return len(self.objectives)

    def get_constraint_count(self) -> int:
        """Get number of constraints."""
        return len(self.constraints)

    def get_design_variable_count(self) -> int:
        """Get number of design variables."""
        return len(self.design_variables)


class BlueprintBuilder:
    """
    Helper class for building StudyBlueprint from interview state.
    """

    def __init__(self):
        """Initialize builder."""
        from .interview_intelligence import InterviewIntelligence
        self.intelligence = InterviewIntelligence()

    def from_interview_state(
        self,
        state: "InterviewState",
        introspection: Optional[Dict[str, Any]] = None
    ) -> StudyBlueprint:
        """
        Build StudyBlueprint from completed interview state.

        Args:
            state: Completed interview state
            introspection: Optional introspection results

        Returns:
            StudyBlueprint ready for generation
        """
        answers = state.answers
        intro = introspection or state.introspection

        # Build design variables
        design_variables = []
        for dv_data in answers.get("design_variables", []):
            if isinstance(dv_data, dict):
                dv = DesignVariable(
                    parameter=dv_data.get("parameter", ""),
                    current_value=dv_data.get("current_value", 0),
                    min_value=dv_data.get("min_value", 0),
                    max_value=dv_data.get("max_value", 1),
                    units=dv_data.get("units"),
                    is_integer=dv_data.get("is_integer", False),
                )
                design_variables.append(dv)
            elif isinstance(dv_data, str):
                # Just a parameter name - look up in introspection
                expr = self._find_expression(dv_data, intro.get("expressions", []))
                if expr:
                    value = expr.get("value", 0)
                    dv = DesignVariable(
                        parameter=dv_data,
                        current_value=value,
                        min_value=value * 0.5 if value > 0 else value * 1.5,
                        max_value=value * 1.5 if value > 0 else value * 0.5,
                    )
                    design_variables.append(dv)

        # Build objectives
        objectives = []
        primary_goal = answers.get("objectives", [{}])
        if isinstance(primary_goal, list) and primary_goal:
            primary = primary_goal[0] if isinstance(primary_goal[0], dict) else {"goal": primary_goal[0]}
        else:
            primary = {"goal": str(primary_goal)}

        # Map to extractor
        extractor_sel = self.intelligence.extractor_mapper.map_goal_to_extractor(
            primary.get("goal", ""),
            intro
        )

        objectives.append(Objective(
            name=primary.get("name", "primary_objective"),
            goal=self._normalize_goal(primary.get("goal", "")),
            extractor=extractor_sel.extractor_id,
            extractor_name=extractor_sel.extractor_name,
            extractor_params=extractor_sel.params,
            weight=primary.get("weight", 1.0),
        ))

        # Add secondary objectives
        secondary = answers.get("objectives_secondary", [])
        for sec_goal in secondary:
            if sec_goal == "none" or not sec_goal:
                continue

            sec_sel = self.intelligence.extractor_mapper.map_goal_to_extractor(
                sec_goal, intro
            )

            objectives.append(Objective(
                name=f"secondary_{sec_goal}",
                goal=self._normalize_goal(sec_goal),
                extractor=sec_sel.extractor_id,
                extractor_name=sec_sel.extractor_name,
                extractor_params=sec_sel.params,
                weight=0.5,  # Default lower weight for secondary
            ))

        # Build constraints
        constraints = []
        constraint_answers = answers.get("constraints", {})
        constraint_handling = answers.get("constraint_handling", "hard")

        if "max_stress" in constraint_answers and constraint_answers["max_stress"]:
            stress_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("stress", intro)
            constraints.append(Constraint(
                name="max_stress",
                constraint_type="max",
                threshold=constraint_answers["max_stress"],
                extractor=stress_sel.extractor_id,
                extractor_name=stress_sel.extractor_name,
                extractor_params=stress_sel.params,
                is_hard=constraint_handling != "soft",
            ))

        if "max_displacement" in constraint_answers and constraint_answers["max_displacement"]:
            disp_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("displacement", intro)
            constraints.append(Constraint(
                name="max_displacement",
                constraint_type="max",
                threshold=constraint_answers["max_displacement"],
                extractor=disp_sel.extractor_id,
                extractor_name=disp_sel.extractor_name,
                extractor_params=disp_sel.params,
                is_hard=constraint_handling != "soft",
            ))

        if "min_frequency" in constraint_answers and constraint_answers["min_frequency"]:
            freq_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("frequency", intro)
            constraints.append(Constraint(
                name="min_frequency",
                constraint_type="min",
                threshold=constraint_answers["min_frequency"],
                extractor=freq_sel.extractor_id,
                extractor_name=freq_sel.extractor_name,
                extractor_params=freq_sel.params,
                is_hard=constraint_handling != "soft",
            ))

        if "max_mass" in constraint_answers and constraint_answers["max_mass"]:
            mass_sel = self.intelligence.extractor_mapper.map_constraint_to_extractor("mass", intro)
            constraints.append(Constraint(
                name="max_mass",
                constraint_type="max",
                threshold=constraint_answers["max_mass"],
                extractor=mass_sel.extractor_id,
                extractor_name=mass_sel.extractor_name,
                is_hard=constraint_handling != "soft",
            ))

        # Determine protocol
        protocol = "protocol_11_multi" if len(objectives) > 1 else "protocol_10_single"

        # Get settings
        n_trials = answers.get("n_trials", 100)
        if n_trials == "custom":
            n_trials = 100  # Default

        # Build blueprint
        blueprint = StudyBlueprint(
            study_name=state.study_name,
            study_description=answers.get("problem_description", ""),
            interview_session_id=state.session_id,
            model_path=intro.get("part_file", ""),
            sim_path=intro.get("sim_file", ""),
            fem_path=intro.get("fem_file", ""),
            design_variables=design_variables,
            objectives=objectives,
            constraints=constraints,
            protocol=protocol,
            n_trials=int(n_trials) if isinstance(n_trials, (int, float)) else 100,
            sampler=self.intelligence.suggest_sampler(len(objectives), len(design_variables)),
            use_neural_acceleration=answers.get("use_neural_acceleration", False),
            solve_all_solutions=answers.get("solve_all_solutions", True),
            warnings_acknowledged=state.warnings_acknowledged,
            baseline_validated=answers.get("run_baseline_validation", False),
        )

        return blueprint

    def _find_expression(self, name: str, expressions: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """Find expression by name."""
        for expr in expressions:
            if expr.get("name") == name:
                return expr
        return None

    def _normalize_goal(self, goal: str) -> str:
        """Normalize goal string to standard format."""
        goal_lower = goal.lower()

        if "minimize" in goal_lower or "reduce" in goal_lower:
            return "minimize"
        elif "maximize" in goal_lower or "increase" in goal_lower:
            return "maximize"
        elif "target" in goal_lower:
            return "target"
        else:
            return goal


# Import for type hints
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from .interview_state import InterviewState