Atomizer/tests/test_migrator.py

"""
Unit tests for SpecMigrator

Tests for migrating legacy optimization_config.json to AtomizerSpec v2.0.

P4.6: Migration tests
"""

import json
import pytest
import tempfile
from pathlib import Path

import sys
sys.path.insert(0, str(Path(__file__).parent.parent))

from optimization_engine.config.migrator import SpecMigrator, MigrationError


# ============================================================================
# Fixtures - Legacy Config Formats
# ============================================================================

@pytest.fixture
def mirror_config() -> dict:
    """Legacy mirror/Zernike config format."""
    return {
        "study_name": "m1_mirror_test",
        "description": "Test mirror optimization",
        "nx_settings": {
            "sim_file": "model.sim",
            "nx_install_path": "C:\\Program Files\\Siemens\\NX2506",
            "simulation_timeout_s": 600
        },
        "zernike_settings": {
            "inner_radius": 100,
            "outer_radius": 500,
            "n_modes": 40,
            "filter_low_orders": 4,
            "displacement_unit": "mm",
            "reference_subcase": 1
        },
        "design_variables": [
            {
                "name": "thickness",
                "parameter": "thickness",
                "bounds": [5.0, 15.0],
                "baseline": 10.0,
                "units": "mm"
            },
            {
                "name": "rib_angle",
                "parameter": "rib_angle",
                "bounds": [20.0, 40.0],
                "baseline": 30.0,
                "units": "degrees"
            }
        ],
        "objectives": [
            {"name": "wfe_40_20", "goal": "minimize", "weight": 10.0},
            {"name": "wfe_mfg", "goal": "minimize", "weight": 1.0},
            {"name": "mass_kg", "goal": "minimize", "weight": 1.0}
        ],
        "constraints": [
            {"name": "mass_limit", "type": "<=", "value": 100.0}
        ],
        "optimization": {
            "algorithm": "TPE",
            "n_trials": 50,
            "seed": 42
        }
    }


@pytest.fixture
def structural_config() -> dict:
    """Legacy structural/bracket config format."""
    return {
        "study_name": "bracket_test",
        "description": "Test bracket optimization",
        "simulation_settings": {
            "sim_file": "bracket.sim",
            "model_file": "bracket.prt",
            "solver": "nastran",
            "solution_type": "SOL101"
        },
        "extraction_settings": {
            "type": "displacement",
            "node_id": 1000,
            "component": "magnitude"
        },
        "design_variables": [
            {
                "name": "thickness",
                "expression_name": "web_thickness",
                "min": 2.0,
                "max": 10.0,
                "baseline": 5.0,
                "units": "mm"
            }
        ],
        "objectives": [
            {"name": "displacement", "type": "minimize", "weight": 1.0},
            {"name": "mass", "direction": "minimize", "weight": 1.0}
        ],
        "constraints": [
            {"name": "stress_limit", "type": "<=", "value": 200.0}
        ],
        "optimization_settings": {
            "sampler": "CMA-ES",
            "n_trials": 100,
            "sigma0": 0.3
        }
    }


@pytest.fixture
def minimal_legacy_config() -> dict:
    """Minimal legacy config for edge case testing."""
    return {
        "study_name": "minimal",
        "design_variables": [
            {"name": "x", "bounds": [0, 1]}
        ],
        "objectives": [
            {"name": "y", "goal": "minimize"}
        ]
    }


# ============================================================================
# Migration Tests
# ============================================================================

class TestSpecMigrator:
    """Tests for SpecMigrator."""

    def test_migrate_mirror_config(self, mirror_config):
        """Test migration of mirror/Zernike config."""
        migrator = SpecMigrator()
        spec = migrator.migrate(mirror_config)

        # Check meta
        assert spec["meta"]["version"] == "2.0"
        assert spec["meta"]["study_name"] == "m1_mirror_test"
        assert "mirror" in spec["meta"]["tags"]

        # Check model
        assert spec["model"]["sim"]["path"] == "model.sim"

        # Check design variables
        assert len(spec["design_variables"]) == 2
        dv = spec["design_variables"][0]
        assert dv["bounds"]["min"] == 5.0
        assert dv["bounds"]["max"] == 15.0
        assert dv["expression_name"] == "thickness"

        # Check extractors
        assert len(spec["extractors"]) >= 1
        ext = spec["extractors"][0]
        assert ext["type"] == "zernike_opd"
        assert ext["config"]["outer_radius_mm"] == 500

        # Check objectives
        assert len(spec["objectives"]) == 3
        obj = spec["objectives"][0]
        assert obj["direction"] == "minimize"

        # Check optimization
        assert spec["optimization"]["algorithm"]["type"] == "TPE"
        assert spec["optimization"]["budget"]["max_trials"] == 50

    def test_migrate_structural_config(self, structural_config):
        """Test migration of structural/bracket config."""
        migrator = SpecMigrator()
        spec = migrator.migrate(structural_config)

        # Check meta
        assert spec["meta"]["version"] == "2.0"

        # Check model
        assert spec["model"]["sim"]["path"] == "bracket.sim"
        assert spec["model"]["sim"]["solver"] == "nastran"

        # Check design variables
        assert len(spec["design_variables"]) == 1
        dv = spec["design_variables"][0]
        assert dv["expression_name"] == "web_thickness"
        assert dv["bounds"]["min"] == 2.0
        assert dv["bounds"]["max"] == 10.0

        # Check optimization
        assert spec["optimization"]["algorithm"]["type"] == "CMA-ES"
        assert spec["optimization"]["algorithm"]["config"]["sigma0"] == 0.3

    def test_migrate_minimal_config(self, minimal_legacy_config):
        """Test migration handles minimal configs."""
        migrator = SpecMigrator()
        spec = migrator.migrate(minimal_legacy_config)

        assert spec["meta"]["study_name"] == "minimal"
        assert len(spec["design_variables"]) == 1
        assert spec["design_variables"][0]["bounds"]["min"] == 0
        assert spec["design_variables"][0]["bounds"]["max"] == 1

    def test_bounds_normalization(self):
        """Test bounds array to object conversion."""
        config = {
            "study_name": "bounds_test",
            "design_variables": [
                {"name": "a", "bounds": [1.0, 5.0]},  # Array format
                {"name": "b", "bounds": {"min": 2.0, "max": 6.0}},  # Object format
                {"name": "c", "min": 3.0, "max": 7.0}  # Separate fields
            ],
            "objectives": [{"name": "y", "goal": "minimize"}]
        }
        migrator = SpecMigrator()
        spec = migrator.migrate(config)

        assert spec["design_variables"][0]["bounds"] == {"min": 1.0, "max": 5.0}
        assert spec["design_variables"][1]["bounds"] == {"min": 2.0, "max": 6.0}
        assert spec["design_variables"][2]["bounds"] == {"min": 3.0, "max": 7.0}

    def test_degenerate_bounds_fixed(self):
        """Test that min >= max is fixed."""
        config = {
            "study_name": "degenerate",
            "design_variables": [
                {"name": "zero", "bounds": [0.0, 0.0]},
                {"name": "reverse", "bounds": [10.0, 5.0]}
            ],
            "objectives": [{"name": "y", "goal": "minimize"}]
        }
        migrator = SpecMigrator()
        spec = migrator.migrate(config)

        # Zero bounds should be expanded
        dv0 = spec["design_variables"][0]
        assert dv0["bounds"]["min"] < dv0["bounds"]["max"]

        # Reversed bounds should be expanded around min
        dv1 = spec["design_variables"][1]
        assert dv1["bounds"]["min"] < dv1["bounds"]["max"]

    def test_algorithm_normalization(self):
        """Test algorithm name normalization."""
        test_cases = [
            ("tpe", "TPE"),
            ("TPESampler", "TPE"),
            ("cma-es", "CMA-ES"),
            ("NSGA-II", "NSGA-II"),
            ("random", "RandomSearch"),
            ("turbo", "SAT_v3"),
            ("unknown_algo", "TPE"),  # Falls back to TPE
        ]

        for old_algo, expected in test_cases:
            config = {
                "study_name": f"algo_test_{old_algo}",
                "design_variables": [{"name": "x", "bounds": [0, 1]}],
                "objectives": [{"name": "y", "goal": "minimize"}],
                "optimization": {"algorithm": old_algo}
            }
            migrator = SpecMigrator()
            spec = migrator.migrate(config)
            assert spec["optimization"]["algorithm"]["type"] == expected, f"Failed for {old_algo}"

    def test_objective_direction_normalization(self):
        """Test objective direction normalization."""
        config = {
            "study_name": "direction_test",
            "design_variables": [{"name": "x", "bounds": [0, 1]}],
            "objectives": [
                {"name": "a", "goal": "minimize"},
                {"name": "b", "type": "maximize"},
                {"name": "c", "direction": "minimize"},
                {"name": "d"}  # No direction - should default
            ]
        }
        migrator = SpecMigrator()
        spec = migrator.migrate(config)

        assert spec["objectives"][0]["direction"] == "minimize"
        assert spec["objectives"][1]["direction"] == "maximize"
        assert spec["objectives"][2]["direction"] == "minimize"
        assert spec["objectives"][3]["direction"] == "minimize"  # Default

    def test_canvas_edges_generated(self, mirror_config):
        """Test that canvas edges are auto-generated."""
        migrator = SpecMigrator()
        spec = migrator.migrate(mirror_config)

        assert "canvas" in spec
        assert "edges" in spec["canvas"]
        assert len(spec["canvas"]["edges"]) > 0

    def test_canvas_positions_assigned(self, mirror_config):
        """Test that canvas positions are assigned to all nodes."""
        migrator = SpecMigrator()
        spec = migrator.migrate(mirror_config)

        # Design variables should have positions
        for dv in spec["design_variables"]:
            assert "canvas_position" in dv
            assert "x" in dv["canvas_position"]
            assert "y" in dv["canvas_position"]

        # Extractors should have positions
        for ext in spec["extractors"]:
            assert "canvas_position" in ext

        # Objectives should have positions
        for obj in spec["objectives"]:
            assert "canvas_position" in obj


class TestMigrationFile:
    """Tests for file-based migration."""

    def test_migrate_file(self, mirror_config):
        """Test migrating from file."""
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create legacy config file
            config_path = Path(tmpdir) / "optimization_config.json"
            with open(config_path, "w") as f:
                json.dump(mirror_config, f)

            # Migrate
            migrator = SpecMigrator(Path(tmpdir))
            spec = migrator.migrate_file(config_path)

            assert spec["meta"]["study_name"] == "m1_mirror_test"

    def test_migrate_file_and_save(self, mirror_config):
        """Test migrating and saving to file."""
        with tempfile.TemporaryDirectory() as tmpdir:
            config_path = Path(tmpdir) / "optimization_config.json"
            output_path = Path(tmpdir) / "atomizer_spec.json"

            with open(config_path, "w") as f:
                json.dump(mirror_config, f)

            migrator = SpecMigrator(Path(tmpdir))
            spec = migrator.migrate_file(config_path, output_path)

            # Check output file was created
            assert output_path.exists()

            # Check content
            with open(output_path) as f:
                saved_spec = json.load(f)
            assert saved_spec["meta"]["version"] == "2.0"

    def test_migrate_file_not_found(self):
        """Test error on missing file."""
        migrator = SpecMigrator()
        with pytest.raises(MigrationError):
            migrator.migrate_file(Path("nonexistent.json"))


# ============================================================================
# Run Tests
# ============================================================================

if __name__ == "__main__":
    pytest.main([__file__, "-v"])