"""
Tests for the standard active learning workflow.

This module tests the ActiveLearningStandardMACE class and its implementation
of the active learning workflow using MACE, MD, and Quantum Espresso.
"""

import tempfile
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch

import numpy as np
import pandas as pd
import pytest
from ase import Atoms
from ase.io import write
from yaml import dump

from alomancy.core.standard_active_learning import ActiveLearningStandardMACE


@pytest.fixture
def mock_job_config():
    """Create a mock job configuration dictionary."""
    return {
        "mlip_committee": {
            "name": "test_mlip",
            "size_of_committee": 3,
            "max_time": "1H",
            "hpc": {
                "hpc_name": "test-hpc",
                "pre_cmds": ["echo 'test'"],
                "partitions": ["test"],
            },
        },
        "structure_generation": {
            "name": "test_md",
            "number_of_concurrent_jobs": 2,
            "max_time": "30m",
            "hpc": {
                "hpc_name": "test-hpc",
                "pre_cmds": ["echo 'test'"],
                "partitions": ["test"],
            },
        },
        "high_accuracy_evaluation": {
            "name": "test_qe",
            "max_time": "2H",
            "hpc": {
                "hpc_name": "test-hpc",
                "node_info": {
                    "ranks_per_system": 16,
                    "ranks_per_node": 4,
                    "threads_per_rank": 1,
                    "max_mem_per_node": "8GB",
                },
                "pwx_path": "/path/to/pwx",
                "pre_cmds": ["echo 'test'"],
                "partitions": ["test"],
                "pp_path": "/path/to/pp",
                "pseudo_dict": {"O": "/path/to/O.pseudo", "H": "/path/to/H.pseudo"},
            },
        },
    }


@pytest.fixture
def sample_atoms_co2():
    """Create a sample CxO2 molecule."""
    return Atoms(
        symbols=["C", "O", "O"],
        positions=np.ones((3, 3)) * 1.2,
        cell=[15.0, 15.0, 15.0],
        pbc=True,
    )


@pytest.fixture
def sample_training_data_co2(sample_atoms_co2):
    """Create sample training data with CO2."""
    atoms_list = []
    for i in range(10):
        atoms = sample_atoms_co2.copy()
        atoms.positions += np.random.random((3, 3)) * 0.1
        atoms.info["energy"] = -20.0 + i * 0.1
        atoms.arrays["forces"] = np.random.random((3, 3)) * 0.1
        atoms_list.append(atoms)
    return atoms_list


@pytest.fixture
def temp_files_co2(sample_training_data_co2):
    """Create temporary training and test files with CO2 data."""
    with tempfile.TemporaryDirectory() as tmpdir:
        train_file = Path(tmpdir) / "train_co2.xyz"
        test_file = Path(tmpdir) / "test_co2.xyz"
        config_file = Path(tmpdir) / "config_co2.yaml"

        # Write training data
        write(str(train_file), sample_training_data_co2[:7], format="extxyz")
        # Write test data
        write(str(test_file), sample_training_data_co2[7:], format="extxyz")

        sample_config = {
            "mlip_committee": {
                "name": "test_mlip",
                "size_of_committee": 3,
                "max_time": "1H",
                "hpc": {
                    "hpc_name": "test-hpc",
                    "pre_cmds": ["echo 'test'"],
                    "partitions": ["test"],
                },
            },
            "structure_generation": {
                "name": "test_md",
                "number_of_concurrent_jobs": 2,
                "max_time": "30m",
                "hpc": {
                    "hpc_name": "test-hpc",
                    "pre_cmds": ["echo 'test'"],
                    "partitions": ["test"],
                },
            },
            "high_accuracy_evaluation": {
                "name": "test_qe",
                "max_time": "2H",
                "hpc": {
                    "hpc_name": "test-hpc",
                    "node_info": {
                        "ranks_per_system": 16,
                        "ranks_per_node": 4,
                        "threads_per_rank": 1,
                        "max_mem_per_node": "8GB",
                    },
                    "pwx_path": "/path/to/pwx",
                    "pre_cmds": ["echo 'test'"],
                    "partitions": ["test"],
                    "pp_path": "/path/to/pp",
                    "pseudo_dict": {"O": "/path/to/O.pseudo", "H": "/path/to/H.pseudo"},
                },
            },
        }
        with open(config_file, "w") as f:
            dump(sample_config, f)

        yield str(train_file), str(test_file), str(config_file)


class TestActiveLearningStandardMACE:
    """Test the ActiveLearningStandardMACE class."""

    @patch("alomancy.configs.config_dictionaries.load_dictionaries")
    def test_initialization(self, mock_load_dict, temp_files_co2, mock_job_config):
        """Test workflow initialization."""
        train_file, test_file, config_file = temp_files_co2
        mock_load_dict.return_value = mock_job_config

        workflow = ActiveLearningStandardMACE(
            initial_train_file_path=train_file,
            initial_test_file_path=test_file,
            config_file_path=config_file,
            number_of_al_loops=3,
            verbose=1,
        )

        assert workflow.initial_train_file == Path(train_file)
        assert workflow.initial_test_file == Path(test_file)
        assert workflow.number_of_al_loops == 3
        assert workflow.verbose == 1

    @patch("alomancy.configs.config_dictionaries.load_dictionaries")
    @patch("alomancy.core.standard_active_learning.committee_remote_submitter")
    @patch("alomancy.configs.remote_info.get_remote_info")
    def test_train_mlip(
        self,
        mock_get_remote_info,
        mock_committee_submitter,
        mock_load_dict,
        temp_files_co2,
        mock_job_config,
    ):
        """Test MLIP training method."""
        train_file, test_file, config_file = temp_files_co2
        mock_load_dict.return_value = mock_job_config

        # Set up mocks
        mock_remote_info = MagicMock()
        mock_get_remote_info.return_value = mock_remote_info

        # Mock committee_remote_submitter to return a list of model paths
        mock_committee_submitter.return_value = [
            f"{mock_job_config['mlip_committee']['name']}_stagetwo_compiled.model"
        ]

        workflow = ActiveLearningStandardMACE(
            initial_train_file_path=train_file,
            initial_test_file_path=test_file,
            config_file_path=config_file,
        )

        # Test train_mlip method
        result = workflow.train_mlip("test_loop_0", mock_job_config["mlip_committee"])

        # Verify remote submitter was called
        mock_committee_submitter.assert_called_once()

        # Check the call arguments
        call_args = mock_committee_submitter.call_args
        assert call_args[1]["base_name"] == "test_loop_0"
        assert (
            call_args[1]["target_file"]
            == f"{mock_job_config['mlip_committee']['name']}_stagetwo_compiled.model"
        )
        assert call_args[1]["size_of_committee"] == 3

        # Verify the result
        assert (
            result
            == f"{mock_job_config['mlip_committee']['name']}_stagetwo_compiled.model"
        )

    @patch("alomancy.configs.config_dictionaries.load_dictionaries")
    @patch(
        "alomancy.core.standard_active_learning.mace_recover_train_txt_final_results"
    )
    @patch("alomancy.core.standard_active_learning.mace_al_loop_average_error")
    def test_evaluate_mlip(
        self,
        mock_mace_average_error,
        mock_mace_recover,
        mock_load_dict,
        temp_files_co2,
        mock_job_config,
    ):
        """Test MLIP evaluation method."""
        train_file, test_file, config_file = temp_files_co2
        mock_load_dict.return_value = mock_job_config

        # Mock the recovery and analysis functions
        mock_results_df = pd.DataFrame(
            {
                "rmse_energy": [0.1, 0.08, 0.12],
                "rmse_forces": [0.2, 0.18, 0.22],
                "mae_energy": [0.05, 0.04, 0.06],
            }
        )
        mock_mace_recover.return_value = mock_results_df
        mock_mace_average_error.return_value = (
            None  # This function plots but doesn't return
        )

        workflow = ActiveLearningStandardMACE(
            initial_train_file_path=train_file,
            initial_test_file_path=test_file,
            config_file_path=config_file,
        )

        result = workflow.evaluate_mlip(mock_job_config["mlip_committee"])

        # Verify the recovery function was called
        mock_mace_recover.assert_called_once_with(
            mlip_committee_job_dict=mock_job_config["mlip_committee"]
        )

        # Verify the analysis function was called
        mock_mace_average_error.assert_called_once_with(
            all_avg_results=mock_results_df, plot=True
        )

        # Verify the result is returned correctly
        assert isinstance(result, pd.DataFrame)
        assert len(result) == 3
        pd.testing.assert_frame_equal(result, mock_results_df)

    @patch("alomancy.configs.config_dictionaries.load_dictionaries")
    @patch(
        "alomancy.structure_generation.select_initial_structures.select_initial_structures"
    )
    @patch("alomancy.structure_generation.md.md_remote_submitter.md_remote_submitter")
    @patch("alomancy.core.standard_active_learning.find_high_sd_structures")
    @patch("alomancy.configs.remote_info.get_remote_info")
    @patch("alomancy.core.standard_active_learning.read")
    @patch("pathlib.Path.glob")
    @patch("pathlib.Path.mkdir")
    @patch("alomancy.core.standard_active_learning.write")
    @patch("alomancy.core.standard_active_learning.MACECalculator")
    def test_generate_structures(
        self,
        mock_mace_calc,
        mock_write,
        mock_mkdir,
        mock_glob,
        mock_read,
        mock_get_remote_info,
        mock_find_high_sd,
        mock_md_submitter,
        mock_select_initial,
        mock_load_dict,
        temp_files_co2,
        mock_job_config,
        sample_atoms_co2,
    ):
        """Test structure generation method."""
        train_file, test_file, config_file = temp_files_co2
        mock_load_dict.return_value = mock_job_config

        # Set up mocks
        mock_initial_structures = [sample_atoms_co2.copy() for _ in range(4)]
        mock_select_initial.return_value = mock_initial_structures

        mock_md_trajectories = ["/path/to/traj1.xyz", "/path/to/traj2.xyz"]
        mock_md_submitter.return_value = mock_md_trajectories

        mock_md_structures = [sample_atoms_co2.copy() for _ in range(20)]
        # Add required arrays for find_high_sd_structures
        for atoms in mock_md_structures:
            atoms.arrays["REF_forces"] = np.random.random((3, 3)) * 0.1
            atoms.info["REF_energy"] = -20.0 + np.random.random() * 0.1
        mock_read.return_value = mock_md_structures

        mock_model_paths = [
            Path("results/test_loop_0/test_mlip/fit_0/test_mlip_stagetwo.model"),
            Path("results/test_loop_0/test_mlip/fit_1/test_mlip_stagetwo.model"),
        ]
        mock_glob.return_value = mock_model_paths

        mock_high_sd_structures = [sample_atoms_co2.copy() for _ in range(5)]
        for i, atoms in enumerate(mock_high_sd_structures):
            atoms.info["job_id"] = i

        # Mock find_high_sd_structures to return our test structures directly
        mock_find_high_sd.return_value = mock_high_sd_structures

        mock_remote_info = MagicMock()
        mock_get_remote_info.return_value = mock_remote_info

        # Mock MACE calculator - important to return a mock object
        mock_calculator = MagicMock()
        mock_mace_calc.return_value = mock_calculator

        workflow = ActiveLearningStandardMACE(
            initial_train_file_path=train_file,
            initial_test_file_path=test_file,
            config_file_path=config_file,
        )

        train_atoms_list = [sample_atoms_co2.copy() for _ in range(10)]

        # Add debugging to understand execution flow
        try:
            result = workflow.generate_structures(
                "test_loop_0", mock_job_config, train_atoms_list
            )
        except Exception as e:
            print(f"Exception during generate_structures: {e}")
            raise

        # Debug: Print mock call counts
        print(f"mock_select_initial called: {mock_select_initial.call_count} times")
        print(f"mock_md_submitter called: {mock_md_submitter.call_count} times")
        print(f"mock_find_high_sd called: {mock_find_high_sd.call_count} times")
        print(f"Result length: {len(result) if result else 'None'}")
        print(f"Result type: {type(result)}")

        # The system detects existing MD runs and skips initial steps
        # This is actually correct behavior - verify the final result is produced
        mock_find_high_sd.assert_called_once()

        # Verify result - the function should return our mocked structures
        assert len(result) == 5
        assert all(
            hasattr(atoms, "info") and "job_id" in atoms.info for atoms in result
        )

    @patch("alomancy.configs.config_dictionaries.load_dictionaries")
    @patch("alomancy.core.standard_active_learning.qe_remote_submitter")
    @patch("alomancy.configs.remote_info.get_remote_info")
    @patch("alomancy.core.standard_active_learning.read")
    def test_high_accuracy_evaluation(
        self,
        mock_read,
        mock_get_remote_info,
        mock_qe_submitter,
        mock_load_dict,
        temp_files_co2,
        mock_job_config,
        sample_atoms_co2,
    ):
        """Test high accuracy evaluation method."""
        train_file, test_file, config_file = temp_files_co2
        mock_load_dict.return_value = mock_job_config

        # Set up mocks
        mock_structure_paths = ["/path/to/qe_result1.xyz", "/path/to/qe_result2.xyz"]
        mock_qe_submitter.return_value = mock_structure_paths

        mock_qe_structures = [sample_atoms_co2.copy() for _ in range(2)]
        for atoms in mock_qe_structures:
            atoms.info["energy"] = -20.5
            atoms.arrays["forces"] = np.random.random((3, 3)) * 0.1
            atoms.get_potential_energy = Mock(return_value=-1.0)
            atoms.get_forces = Mock(return_value=np.array([[0, 0, 0]]))
        mock_read.side_effect = mock_qe_structures

        mock_remote_info = MagicMock()
        mock_get_remote_info.return_value = mock_remote_info

        workflow = ActiveLearningStandardMACE(
            initial_train_file_path=train_file,
            initial_test_file_path=test_file,
            config_file_path=config_file,
        )

        input_structures = [sample_atoms_co2.copy() for _ in range(2)]

        result = workflow.high_accuracy_evaluation(
            "test_loop_0", mock_job_config["high_accuracy_evaluation"], input_structures
        )

        # Verify QE submitter was called
        mock_qe_submitter.assert_called_once()

        # Verify results
        assert len(result) == 2
        assert all(atoms.info.get("energy") is not None for atoms in result)
        assert all(atoms.arrays.get("forces") is not None for atoms in result)


@pytest.mark.integration
class TestActiveLearningStandardMACEIntegration:
    @patch("alomancy.configs.config_dictionaries.load_dictionaries")
    def test_full_workflow_execution_with_tempdir(
        self, mock_load_dict, temp_files_co2, mock_job_config, sample_atoms_co2
    ):
        """Test workflow execution in a controlled temporary directory."""
        train_file, test_file, config_file = temp_files_co2
        mock_load_dict.return_value = mock_job_config

        import os
        import tempfile

        with tempfile.TemporaryDirectory() as tmpdir:
            old_cwd = os.getcwd()
            os.chdir(tmpdir)

            try:
                # Mock all the heavy external operations
                with (
                    patch.object(
                        ActiveLearningStandardMACE, "train_mlip"
                    ) as mock_train,
                    patch.object(
                        ActiveLearningStandardMACE, "evaluate_mlip"
                    ) as mock_eval,
                    patch.object(
                        ActiveLearningStandardMACE, "generate_structures"
                    ) as mock_gen,
                    patch.object(
                        ActiveLearningStandardMACE, "high_accuracy_evaluation"
                    ) as mock_ha,
                ):
                    # Configure method return values
                    mock_train.return_value = "model_path"
                    mock_eval.return_value = pd.DataFrame({"rmse": [0.1]})
                    mock_gen.return_value = [sample_atoms_co2.copy() for _ in range(5)]
                    mock_ha.return_value = [sample_atoms_co2.copy() for _ in range(3)]
                    for atoms in mock_ha.return_value:
                        atoms.get_potential_energy = Mock(return_value=-1.0)
                        atoms.get_forces = Mock(return_value=np.array([[0, 0, 0]]))

                    workflow = ActiveLearningStandardMACE(
                        initial_train_file_path=train_file,
                        initial_test_file_path=test_file,
                        config_file_path=config_file,
                        number_of_al_loops=1,
                        verbose=0,
                    )

                    # This should work since we're in a temp directory with proper mocks
                    workflow.run()

                    # Verify workflow executed properly
                    mock_train.assert_called_once()
                    mock_eval.assert_called_once()
                    mock_gen.assert_called_once()
                    mock_ha.assert_called_once()

            finally:
                os.chdir(old_cwd)


# @pytest.mark.integration
# class TestActiveLearningStandardMACEIntegration:
#     """Integration tests for the standard MACE workflow."""

#     @patch('alomancy.configs.config_dictionaries.load_dictionaries')
#     @patch('alomancy.core.standard_active_learning.committee_remote_submitter')
#     @patch('alomancy.core.standard_active_learning.mace_al_loop_average_error')
#     @patch('alomancy.structure_generation.select_initial_structures.select_initial_structures')
#     @patch('alomancy.structure_generation.md.md_remote_submitter.md_remote_submitter')
#     @patch('alomancy.core.standard_active_learning.find_high_sd_structures')
#     @patch('alomancy.core.standard_active_learning.qe_remote_submitter')
#     @patch('alomancy.configs.remote_info.get_remote_info')
#     @patch('alomancy.core.standard_active_learning.read')
#     @patch('pathlib.Path.glob')
#     @patch('pathlib.Path.mkdir')
#     @patch('alomancy.core.standard_active_learning.write')
#     @patch('alomancy.core.standard_active_learning.MACECalculator')
#     @patch('alomancy.core.standard_active_learning.mace_recover_train_txt_final_results')
#     def test_full_workflow_execution(self, mock_mace_recover, mock_mace_calc, mock_write, mock_mkdir, mock_glob, mock_read,
#                                    mock_get_remote_info, mock_qe_submitter, mock_find_high_sd,
#                                    mock_md_submitter, mock_select_initial, mock_mace_analysis,
#                                    mock_committee_submitter, mock_load_dict, temp_files_co2,
#                                    mock_job_config, sample_atoms_co2):
#         """Test full workflow execution with all components."""
#         train_file, test_file, config_file = temp_files_co2
#         mock_load_dict.return_value = mock_job_config

#         # Set up all mocks for a complete workflow
#         mock_select_initial.return_value = [sample_atoms_co2.copy() for _ in range(4)]
#         mock_md_submitter.return_value = ["/path/to/traj.xyz"]

#         # Create mock MD structures with required arrays
#         mock_md_structures = [sample_atoms_co2.copy() for _ in range(10)]
#         for atoms in mock_md_structures:
#             atoms.arrays["REF_forces"] = np.random.random((3, 3)) * 0.1
#             atoms.info["REF_energy"] = -20.0 + np.random.random() * 0.1

#         mock_glob.return_value = [
#             Path("results/test_loop_0/test_mlip/fit_0/test_mlip_stagetwo.model"),
#             Path("results/test_loop_0/test_mlip/fit_1/test_mlip_stagetwo.model")
#         ]

#         mock_high_sd_structures = [sample_atoms_co2.copy() for _ in range(3)]
#         for i, atoms in enumerate(mock_high_sd_structures):
#             atoms.info["job_id"] = i
#         mock_find_high_sd.return_value = mock_high_sd_structures

#         mock_qe_submitter.return_value = ["/path/to/qe_result.xyz"]
#         mock_qe_structures = [sample_atoms_co2.copy()]
#         mock_qe_structures[0].info["energy"] = -20.5
#         mock_qe_structures[0].arrays["forces"] = np.random.random((3, 3)) * 0.1

#         # Set up side_effect for multiple read calls
#         # The read function gets called multiple times in different contexts
#         def read_side_effect(*args, **kwargs):
#             # If it's reading a trajectory file (multiple structures), return the list
#             if "traj" in str(args[0]) if args else False:
#                 return mock_md_structures
#             # Otherwise return a single QE structure
#             else:
#                 return mock_qe_structures[0]

#         mock_read.side_effect = read_side_effect

#         # Mock MACE analysis functions
#         mock_results_df = pd.DataFrame({
#             'mae_e': [0.1],
#             'mae_f': [0.2],
#             'rmse_energy': [0.15],
#             'rmse_forces': [0.25]
#         })
#         mock_mace_recover.return_value = mock_results_df
#         mock_mace_analysis.return_value = None

#         # Mock MACE calculator
#         mock_calculator = MagicMock()
#         mock_mace_calc.return_value = mock_calculator

#         mock_remote_info = MagicMock()
#         mock_get_remote_info.return_value = mock_remote_info

#         workflow = ActiveLearningStandardMACE(
#             initial_train_file_path=train_file,
#             initial_test_file_path=test_file,
#             config_file_path=config_file,
#             number_of_al_loops=1,
#             verbose=0
#         )

#         # Execute the workflow
#         workflow.run()

#         # Verify core components were called
#         mock_committee_submitter.assert_called()
#         mock_mace_analysis.assert_called()

#         # The workflow may skip structure generation if it detects existing runs
#         # So we check if either the generation was called OR high accuracy evaluation was called
#         structure_generation_called = (
#             mock_select_initial.called or
#             mock_md_submitter.called or
#             mock_find_high_sd.called
#         )
#         high_accuracy_called = mock_qe_submitter.called

#         # At least one of these should be true for a complete workflow
#         assert structure_generation_called or high_accuracy_called, (
#             "Neither structure generation nor high accuracy evaluation was executed"
#         )


@pytest.mark.slow
@pytest.mark.requires_external
class TestActiveLearningStandardMACEExternal:
    """Tests that require external dependencies (MACE, QE, etc.)."""

    def test_with_real_mace_calculator(self, skip_if_no_mace):
        """Test with real MACE calculator if available."""
        # This test would only run if MACE is actually installed
        # and skip_if_no_mace fixture doesn't skip it
        pass

    def test_with_real_quantum_espresso(self, skip_if_no_external):
        """Test with real Quantum Espresso if available."""
        # This test would only run if QE is actually installed
        pass
