Atomizer/tests/interview/test_study_blueprint.py

"""Tests for StudyBlueprint and BlueprintBuilder."""

import pytest
import json

from optimization_engine.interview.study_blueprint import (
    StudyBlueprint,
    DesignVariable,
    Objective,
    Constraint,
    BlueprintBuilder,
)
from optimization_engine.interview.interview_state import InterviewState


class TestDesignVariable:
    """Tests for DesignVariable dataclass."""

    def test_creation(self):
        """Test creating design variable."""
        dv = DesignVariable(
            parameter="thickness",
            current_value=5.0,
            min_value=1.0,
            max_value=10.0,
            units="mm"
        )

        assert dv.parameter == "thickness"
        assert dv.min_value == 1.0
        assert dv.max_value == 10.0
        assert dv.current_value == 5.0
        assert dv.units == "mm"

    def test_to_dict(self):
        """Test conversion to dict."""
        dv = DesignVariable(
            parameter="thickness",
            current_value=5.0,
            min_value=1.0,
            max_value=10.0
        )

        d = dv.to_dict()
        assert d["parameter"] == "thickness"
        assert d["min_value"] == 1.0
        assert d["max_value"] == 10.0

    def test_to_config_format(self):
        """Test conversion to config format."""
        dv = DesignVariable(
            parameter="thickness",
            current_value=5.0,
            min_value=1.0,
            max_value=10.0,
            units="mm"
        )

        config = dv.to_config_format()
        assert config["expression_name"] == "thickness"
        assert config["bounds"] == [1.0, 10.0]


class TestObjective:
    """Tests for Objective dataclass."""

    def test_creation(self):
        """Test creating objective."""
        obj = Objective(
            name="mass",
            goal="minimize",
            extractor="E4",
            weight=1.0
        )

        assert obj.name == "mass"
        assert obj.goal == "minimize"
        assert obj.extractor == "E4"
        assert obj.weight == 1.0

    def test_to_dict(self):
        """Test conversion to dict."""
        obj = Objective(
            name="displacement",
            goal="minimize",
            extractor="E1",
            extractor_params={"node_id": 123}
        )

        d = obj.to_dict()
        assert d["name"] == "displacement"
        assert d["extractor"] == "E1"
        assert d["extractor_params"]["node_id"] == 123

    def test_to_config_format(self):
        """Test conversion to config format."""
        obj = Objective(
            name="mass",
            goal="minimize",
            extractor="E4",
            weight=0.5
        )

        config = obj.to_config_format()
        assert config["name"] == "mass"
        assert config["type"] == "minimize"
        assert config["weight"] == 0.5


class TestConstraint:
    """Tests for Constraint dataclass."""

    def test_creation(self):
        """Test creating constraint."""
        con = Constraint(
            name="max_stress",
            constraint_type="max",
            threshold=200.0,
            extractor="E3"
        )

        assert con.name == "max_stress"
        assert con.constraint_type == "max"
        assert con.threshold == 200.0

    def test_to_dict(self):
        """Test conversion to dict."""
        con = Constraint(
            name="max_displacement",
            constraint_type="max",
            threshold=0.5,
            extractor="E1"
        )

        d = con.to_dict()
        assert d["name"] == "max_displacement"
        assert d["threshold"] == 0.5

    def test_to_config_format(self):
        """Test conversion to config format."""
        con = Constraint(
            name="max_stress",
            constraint_type="max",
            threshold=200.0,
            extractor="E3",
            is_hard=True
        )

        config = con.to_config_format()
        assert config["type"] == "max"
        assert config["threshold"] == 200.0
        assert config["hard"] is True


class TestStudyBlueprint:
    """Tests for StudyBlueprint dataclass."""

    def test_creation(self):
        """Test creating blueprint."""
        bp = StudyBlueprint(
            study_name="test_study",
            study_description="A test study",
            model_path="/path/model.prt",
            sim_path="/path/sim.sim",
            design_variables=[
                DesignVariable(parameter="t", current_value=5, min_value=1, max_value=10)
            ],
            objectives=[
                Objective(name="mass", goal="minimize", extractor="E4")
            ],
            constraints=[
                Constraint(name="stress", constraint_type="max", threshold=200, extractor="E3")
            ],
            protocol="protocol_10_single",
            n_trials=100,
            sampler="TPE"
        )

        assert bp.study_name == "test_study"
        assert len(bp.design_variables) == 1
        assert len(bp.objectives) == 1
        assert len(bp.constraints) == 1

    def test_to_config_json(self):
        """Test conversion to optimization_config.json format."""
        bp = StudyBlueprint(
            study_name="test",
            study_description="Test",
            model_path="/model.prt",
            sim_path="/sim.sim",
            design_variables=[
                DesignVariable(parameter="thickness", current_value=5, min_value=1, max_value=10)
            ],
            objectives=[
                Objective(name="mass", goal="minimize", extractor="E4")
            ],
            constraints=[],
            protocol="protocol_10_single",
            n_trials=50,
            sampler="TPE"
        )

        config = bp.to_config_json()

        assert isinstance(config, dict)
        assert config["study_name"] == "test"
        assert "design_variables" in config
        assert "objectives" in config

        # Should be valid JSON
        json_str = json.dumps(config)
        assert len(json_str) > 0

    def test_to_markdown(self):
        """Test conversion to markdown summary."""
        bp = StudyBlueprint(
            study_name="bracket_v1",
            study_description="Bracket optimization",
            model_path="/model.prt",
            sim_path="/sim.sim",
            design_variables=[
                DesignVariable(parameter="thickness", current_value=5, min_value=1, max_value=10, units="mm")
            ],
            objectives=[
                Objective(name="mass", goal="minimize", extractor="E4")
            ],
            constraints=[
                Constraint(name="stress", constraint_type="max", threshold=200, extractor="E3")
            ],
            protocol="protocol_10_single",
            n_trials=100,
            sampler="TPE"
        )

        md = bp.to_markdown()

        assert "bracket_v1" in md
        assert "thickness" in md
        assert "mass" in md.lower()
        assert "stress" in md
        assert "200" in md
        assert "100" in md  # n_trials
        assert "TPE" in md

    def test_validate_valid_blueprint(self):
        """Test validation passes for valid blueprint."""
        bp = StudyBlueprint(
            study_name="test",
            study_description="Test",
            model_path="/model.prt",
            sim_path="/sim.sim",
            design_variables=[
                DesignVariable(parameter="t", current_value=5, min_value=1, max_value=10)
            ],
            objectives=[
                Objective(name="mass", goal="minimize", extractor="E4")
            ],
            constraints=[
                Constraint(name="stress", constraint_type="max", threshold=200, extractor="E3")
            ],
            protocol="protocol_10_single",
            n_trials=100,
            sampler="TPE"
        )

        errors = bp.validate()
        assert len(errors) == 0

    def test_validate_missing_objectives(self):
        """Test validation catches missing objectives."""
        bp = StudyBlueprint(
            study_name="test",
            study_description="Test",
            model_path="/model.prt",
            sim_path="/sim.sim",
            design_variables=[
                DesignVariable(parameter="t", current_value=5, min_value=1, max_value=10)
            ],
            objectives=[],  # No objectives
            constraints=[],
            protocol="protocol_10_single",
            n_trials=100,
            sampler="TPE"
        )

        errors = bp.validate()
        assert any("objective" in e.lower() for e in errors)

    def test_validate_missing_design_variables(self):
        """Test validation catches missing design variables."""
        bp = StudyBlueprint(
            study_name="test",
            study_description="Test",
            model_path="/model.prt",
            sim_path="/sim.sim",
            design_variables=[],  # No design variables
            objectives=[
                Objective(name="mass", goal="minimize", extractor="E4")
            ],
            constraints=[],
            protocol="protocol_10_single",
            n_trials=100,
            sampler="TPE"
        )

        errors = bp.validate()
        assert any("design variable" in e.lower() for e in errors)

    def test_validate_invalid_bounds(self):
        """Test validation catches invalid bounds."""
        bp = StudyBlueprint(
            study_name="test",
            study_description="Test",
            model_path="/model.prt",
            sim_path="/sim.sim",
            design_variables=[
                DesignVariable(parameter="t", current_value=5, min_value=10, max_value=1)  # min > max
            ],
            objectives=[
                Objective(name="mass", goal="minimize", extractor="E4")
            ],
            constraints=[],
            protocol="protocol_10_single",
            n_trials=100,
            sampler="TPE"
        )

        errors = bp.validate()
        assert any("bound" in e.lower() or "min" in e.lower() for e in errors)

    def test_to_dict_from_dict_roundtrip(self):
        """Test dict serialization roundtrip."""
        bp = StudyBlueprint(
            study_name="test",
            study_description="Test",
            model_path="/model.prt",
            sim_path="/sim.sim",
            design_variables=[
                DesignVariable(parameter="t", current_value=5, min_value=1, max_value=10)
            ],
            objectives=[
                Objective(name="mass", goal="minimize", extractor="E4")
            ],
            constraints=[
                Constraint(name="stress", constraint_type="max", threshold=200, extractor="E3")
            ],
            protocol="protocol_10_single",
            n_trials=100,
            sampler="TPE"
        )

        d = bp.to_dict()
        bp2 = StudyBlueprint.from_dict(d)

        assert bp2.study_name == bp.study_name
        assert len(bp2.design_variables) == len(bp.design_variables)
        assert bp2.n_trials == bp.n_trials


class TestBlueprintBuilder:
    """Tests for BlueprintBuilder."""

    def test_from_interview_state_simple(self):
        """Test building blueprint from simple interview state."""
        builder = BlueprintBuilder()

        state = InterviewState(
            study_name="bracket_v1",
            study_path="/path/to/study"
        )
        state.answers = {
            "study_description": "Bracket mass optimization",
            "objectives": [{"goal": "minimize_mass"}],
            "constraints": [{"type": "stress", "threshold": 200}],
            "design_variables": [
                {"name": "thickness", "min": 1, "max": 10, "current": 5}
            ],
            "n_trials": 100,
        }

        introspection = {
            "model_path": "/path/model.prt",
            "sim_path": "/path/sim.sim"
        }

        bp = builder.from_interview_state(state, introspection)

        assert bp.study_name == "bracket_v1"
        assert len(bp.design_variables) >= 1
        assert len(bp.objectives) >= 1

    def test_from_interview_state_multi_objective(self):
        """Test building blueprint for multi-objective optimization."""
        builder = BlueprintBuilder()

        state = InterviewState(study_name="multi_obj")
        state.answers = {
            "study_description": "Multi-objective",
            "objectives": [
                {"goal": "minimize_mass"},
                {"goal": "minimize_displacement"}
            ],
            "constraints": [],
            "design_variables": [
                {"name": "t", "min": 1, "max": 10}
            ],
            "n_trials": 200
        }

        introspection = {}

        bp = builder.from_interview_state(state, introspection)

        # Blueprint creation succeeds
        assert bp is not None
        assert bp.study_name == "multi_obj"

    def test_auto_assign_extractors(self):
        """Test automatic extractor assignment."""
        builder = BlueprintBuilder()

        state = InterviewState(study_name="test")
        state.answers = {
            "study_description": "Test",
            "objectives": [{"goal": "minimize_mass"}],  # No extractor specified
            "constraints": [],
            "design_variables": [{"name": "t", "min": 1, "max": 10}],
            "n_trials": 50
        }

        bp = builder.from_interview_state(state, {})

        # Should auto-assign E4 for mass
        assert bp.objectives[0].extractor == "E4"

    def test_calculate_n_trials(self):
        """Test automatic trial count calculation."""
        builder = BlueprintBuilder()

        # Few design variables = fewer trials
        state = InterviewState(study_name="test")
        state.answers = {
            "study_description": "Test",
            "objectives": [{"goal": "minimize_mass"}],
            "constraints": [],
            "design_variables": [
                {"name": "t1", "min": 1, "max": 10},
                {"name": "t2", "min": 1, "max": 10},
            ],
        }
        state.complexity = "simple"

        bp = builder.from_interview_state(state, {})
        assert bp.n_trials >= 50  # Minimum trials

    def test_select_sampler(self):
        """Test automatic sampler selection."""
        builder = BlueprintBuilder()

        # Single objective = TPE
        state = InterviewState(study_name="test")
        state.answers = {
            "study_description": "Test",
            "objectives": [{"goal": "minimize_mass"}],
            "constraints": [],
            "design_variables": [{"name": "t", "min": 1, "max": 10}],
        }

        bp = builder.from_interview_state(state, {})
        assert bp.sampler == "TPE"

        # Multi-objective case - sampler selection depends on implementation
        state.answers["objectives"] = [
            {"goal": "minimize_mass"},
            {"goal": "minimize_displacement"}
        ]

        bp = builder.from_interview_state(state, {})
        # Just verify blueprint is created successfully
        assert bp.sampler is not None