"""
Unit tests for TrainingDataExporter

Tests the training data export functionality for AtomizerField.
"""

import pytest
import json
import tempfile
import shutil
from pathlib import Path
from optimization_engine.processors.surrogates.training_data_exporter import TrainingDataExporter, create_exporter_from_config


class TestTrainingDataExporter:
    """Test suite for TrainingDataExporter class."""

    @pytest.fixture
    def temp_dir(self):
        """Create a temporary directory for tests."""
        temp_path = Path(tempfile.mkdtemp())
        yield temp_path
        # Cleanup after test
        if temp_path.exists():
            shutil.rmtree(temp_path)

    @pytest.fixture
    def sample_exporter(self, temp_dir):
        """Create a sample exporter for testing."""
        return TrainingDataExporter(
            export_dir=temp_dir / "test_export",
            study_name="test_study",
            design_variable_names=["thickness", "width", "length"],
            objective_names=["max_stress", "mass"],
            constraint_names=["stress_constraint"],
            metadata={"test_key": "test_value"}
        )

    @pytest.fixture
    def sample_simulation_files(self, temp_dir):
        """Create sample .dat and .op2 files for testing."""
        sim_dir = temp_dir / "sim_files"
        sim_dir.mkdir(parents=True, exist_ok=True)

        dat_file = sim_dir / "test_model.dat"
        op2_file = sim_dir / "test_model.op2"

        # Create dummy files with some content
        dat_file.write_text("$ Nastran input deck\nGRID,1,0,0.0,0.0,0.0\n")
        op2_file.write_bytes(b"DUMMY OP2 BINARY DATA" * 100)

        return {
            "dat_file": dat_file,
            "op2_file": op2_file
        }

    def test_exporter_initialization(self, temp_dir):
        """Test that exporter initializes correctly."""
        export_dir = temp_dir / "test_export"

        exporter = TrainingDataExporter(
            export_dir=export_dir,
            study_name="beam_study",
            design_variable_names=["thickness", "width"],
            objective_names=["stress", "mass"]
        )

        assert exporter.export_dir == export_dir
        assert exporter.study_name == "beam_study"
        assert exporter.design_variable_names == ["thickness", "width"]
        assert exporter.objective_names == ["stress", "mass"]
        assert exporter.trial_count == 0
        assert len(exporter.exported_trials) == 0
        assert export_dir.exists()

    def test_readme_creation(self, sample_exporter):
        """Test that README.md is created."""
        readme_path = sample_exporter.export_dir / "README.md"
        assert readme_path.exists()

        content = readme_path.read_text()
        assert "test_study" in content
        assert "AtomizerField Training Data" in content
        assert "thickness" in content
        assert "max_stress" in content

    def test_export_trial_success(self, sample_exporter, sample_simulation_files):
        """Test successful trial export."""
        design_vars = {"thickness": 3.5, "width": 50.0, "length": 200.0}
        results = {
            "objectives": {"max_stress": 245.3, "mass": 1.25},
            "constraints": {"stress_constraint": -54.7},
            "max_displacement": 1.23
        }

        success = sample_exporter.export_trial(
            trial_number=1,
            design_variables=design_vars,
            results=results,
            simulation_files=sample_simulation_files
        )

        assert success
        assert sample_exporter.trial_count == 1
        assert len(sample_exporter.exported_trials) == 1

        # Check directory structure
        trial_dir = sample_exporter.export_dir / "trial_0001"
        assert trial_dir.exists()
        assert (trial_dir / "input").exists()
        assert (trial_dir / "output").exists()
        assert (trial_dir / "input" / "model.bdf").exists()
        assert (trial_dir / "output" / "model.op2").exists()
        assert (trial_dir / "metadata.json").exists()

    def test_export_trial_metadata(self, sample_exporter, sample_simulation_files):
        """Test that metadata.json is created correctly."""
        design_vars = {"thickness": 3.5, "width": 50.0, "length": 200.0}
        results = {
            "objectives": {"max_stress": 245.3, "mass": 1.25},
            "constraints": {"stress_constraint": -54.7},
            "max_displacement": 1.23
        }

        sample_exporter.export_trial(
            trial_number=42,
            design_variables=design_vars,
            results=results,
            simulation_files=sample_simulation_files
        )

        metadata_path = sample_exporter.export_dir / "trial_0042" / "metadata.json"
        assert metadata_path.exists()

        with open(metadata_path, 'r') as f:
            metadata = json.load(f)

        assert metadata["trial_number"] == 42
        assert metadata["atomizer_study"] == "test_study"
        assert metadata["design_parameters"] == design_vars
        assert metadata["results"]["objectives"] == {"max_stress": 245.3, "mass": 1.25}
        assert metadata["results"]["constraints"] == {"stress_constraint": -54.7}
        assert metadata["results"]["max_displacement"] == 1.23
        assert "timestamp" in metadata

    def test_export_trial_missing_dat_file(self, sample_exporter, temp_dir):
        """Test export fails gracefully when .dat file is missing."""
        design_vars = {"thickness": 3.5}
        results = {"objectives": {"max_stress": 245.3}}

        # Create only .op2 file
        sim_dir = temp_dir / "sim_files"
        sim_dir.mkdir(parents=True, exist_ok=True)
        op2_file = sim_dir / "test_model.op2"
        op2_file.write_bytes(b"DUMMY OP2 DATA")

        simulation_files = {
            "dat_file": sim_dir / "nonexistent.dat",  # Does not exist
            "op2_file": op2_file
        }

        success = sample_exporter.export_trial(
            trial_number=1,
            design_variables=design_vars,
            results=results,
            simulation_files=simulation_files
        )

        assert not success
        assert sample_exporter.trial_count == 0

    def test_export_trial_missing_op2_file(self, sample_exporter, temp_dir):
        """Test export fails gracefully when .op2 file is missing."""
        design_vars = {"thickness": 3.5}
        results = {"objectives": {"max_stress": 245.3}}

        # Create only .dat file
        sim_dir = temp_dir / "sim_files"
        sim_dir.mkdir(parents=True, exist_ok=True)
        dat_file = sim_dir / "test_model.dat"
        dat_file.write_text("$ Nastran input")

        simulation_files = {
            "dat_file": dat_file,
            "op2_file": sim_dir / "nonexistent.op2"  # Does not exist
        }

        success = sample_exporter.export_trial(
            trial_number=1,
            design_variables=design_vars,
            results=results,
            simulation_files=simulation_files
        )

        assert not success
        assert sample_exporter.trial_count == 0

    def test_export_multiple_trials(self, sample_exporter, sample_simulation_files):
        """Test exporting multiple trials."""
        for i in range(1, 6):
            design_vars = {"thickness": 2.0 + i * 0.5, "width": 40.0 + i * 5, "length": 180.0 + i * 10}
            results = {
                "objectives": {"max_stress": 200.0 + i * 10, "mass": 1.0 + i * 0.1}
            }

            success = sample_exporter.export_trial(
                trial_number=i,
                design_variables=design_vars,
                results=results,
                simulation_files=sample_simulation_files
            )

            assert success

        assert sample_exporter.trial_count == 5
        assert len(sample_exporter.exported_trials) == 5

        # Verify all trial directories exist
        for i in range(1, 6):
            trial_dir = sample_exporter.export_dir / f"trial_{i:04d}"
            assert trial_dir.exists()
            assert (trial_dir / "metadata.json").exists()

    def test_finalize(self, sample_exporter, sample_simulation_files):
        """Test finalize creates study_summary.json."""
        # Export a few trials
        for i in range(1, 4):
            design_vars = {"thickness": 2.0 + i * 0.5, "width": 50.0, "length": 200.0}
            results = {"objectives": {"max_stress": 200.0 + i * 10, "mass": 1.2}}
            sample_exporter.export_trial(i, design_vars, results, sample_simulation_files)

        # Finalize
        sample_exporter.finalize()

        # Check study_summary.json
        summary_path = sample_exporter.export_dir / "study_summary.json"
        assert summary_path.exists()

        with open(summary_path, 'r') as f:
            summary = json.load(f)

        assert summary["study_name"] == "test_study"
        assert summary["total_trials"] == 3
        assert summary["design_variables"] == ["thickness", "width", "length"]
        assert summary["objectives"] == ["max_stress", "mass"]
        assert summary["constraints"] == ["stress_constraint"]
        assert "export_timestamp" in summary
        assert summary["metadata"] == {"test_key": "test_value"}

    def test_create_exporter_from_config_disabled(self):
        """Test that None is returned when export is disabled."""
        config = {
            "training_data_export": {
                "enabled": False,
                "export_dir": "/some/path"
            }
        }

        exporter = create_exporter_from_config(config)
        assert exporter is None

    def test_create_exporter_from_config_missing_export_dir(self):
        """Test that None is returned when export_dir is missing."""
        config = {
            "training_data_export": {
                "enabled": True
            }
        }

        exporter = create_exporter_from_config(config)
        assert exporter is None

    def test_create_exporter_from_config_success(self, temp_dir):
        """Test successful exporter creation from config."""
        config = {
            "study_name": "beam_optimization",
            "training_data_export": {
                "enabled": True,
                "export_dir": str(temp_dir / "export")
            },
            "design_variables": [
                {"name": "thickness"},
                {"parameter": "width"},
                {"other_key": "value"}  # Should use "var_2"
            ],
            "objectives": [
                {"name": "max_stress"},
                {"name": "mass"}
            ],
            "constraints": [
                {"name": "stress_limit"}
            ],
            "version": "1.0",
            "optimization": {
                "algorithm": "NSGA-II",
                "n_trials": 100
            }
        }

        exporter = create_exporter_from_config(config)

        assert exporter is not None
        assert exporter.study_name == "beam_optimization"
        assert exporter.design_variable_names == ["thickness", "width", "var_2"]
        assert exporter.objective_names == ["max_stress", "mass"]
        assert exporter.constraint_names == ["stress_limit"]
        assert exporter.study_metadata["atomizer_version"] == "1.0"
        assert exporter.study_metadata["optimization_algorithm"] == "NSGA-II"
        assert exporter.study_metadata["n_trials"] == 100

    def test_create_exporter_from_config_defaults(self, temp_dir):
        """Test exporter creation with minimal config."""
        config = {
            "training_data_export": {
                "enabled": True,
                "export_dir": str(temp_dir / "export")
            }
        }

        exporter = create_exporter_from_config(config)

        assert exporter is not None
        assert exporter.study_name == "unnamed_study"
        assert exporter.design_variable_names == []
        assert exporter.objective_names == []
        assert exporter.constraint_names == []

    def test_file_content_preservation(self, sample_exporter, temp_dir):
        """Test that file contents are preserved during copy."""
        # Create files with specific content
        sim_dir = temp_dir / "sim_files"
        sim_dir.mkdir(parents=True, exist_ok=True)

        dat_content = "$ Nastran Input Deck\nGRID,1,0,0.0,0.0,0.0\nGRID,2,0,1.0,0.0,0.0\n"
        op2_content = b"\x00\x01\x02\x03NASTRAN_BINARY_DATA\xFF\xFE\xFD\xFC" * 50

        dat_file = sim_dir / "model.dat"
        op2_file = sim_dir / "model.op2"
        dat_file.write_text(dat_content)
        op2_file.write_bytes(op2_content)

        simulation_files = {"dat_file": dat_file, "op2_file": op2_file}

        # Export trial
        sample_exporter.export_trial(
            trial_number=1,
            design_variables={"thickness": 3.0},
            results={"objectives": {"stress": 200.0}},
            simulation_files=simulation_files
        )

        # Verify copied files have same content
        exported_dat = sample_exporter.export_dir / "trial_0001" / "input" / "model.bdf"
        exported_op2 = sample_exporter.export_dir / "trial_0001" / "output" / "model.op2"

        assert exported_dat.read_text() == dat_content
        assert exported_op2.read_bytes() == op2_content

    def test_trial_numbering_format(self, sample_exporter, sample_simulation_files):
        """Test that trial directories use correct numbering format."""
        # Test various trial numbers
        trial_numbers = [1, 9, 10, 99, 100, 999, 1000]

        for trial_num in trial_numbers:
            sample_exporter.export_trial(
                trial_number=trial_num,
                design_variables={"thickness": 3.0},
                results={"objectives": {"stress": 200.0}},
                simulation_files=sample_simulation_files
            )

            expected_dir = sample_exporter.export_dir / f"trial_{trial_num:04d}"
            assert expected_dir.exists()

        # Verify specific formatting
        assert (sample_exporter.export_dir / "trial_0001").exists()
        assert (sample_exporter.export_dir / "trial_0009").exists()
        assert (sample_exporter.export_dir / "trial_0010").exists()
        assert (sample_exporter.export_dir / "trial_0099").exists()
        assert (sample_exporter.export_dir / "trial_0100").exists()
        assert (sample_exporter.export_dir / "trial_0999").exists()
        assert (sample_exporter.export_dir / "trial_1000").exists()


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