import unittest
import textwrap
from unittest.mock import Mock, patch, mock_open
from collections import Counter

from vasppy.calculation import (
    Calculation,
    delta_E,
    delta_stoichiometry,
    energy_string_to_float,
    import_calculations_from_file,
)


class CalculationTestCase(unittest.TestCase):
    def test_calculation_is_initialised(self):
        this_title = "A"
        this_energy = -100.0
        this_stoichiometry = {"B": 1, "C": 2}
        calculation = Calculation(
            title=this_title, energy=this_energy, stoichiometry=this_stoichiometry
        )
        self.assertEqual(this_title, calculation.title)
        self.assertEqual(this_energy, calculation.energy)
        self.assertEqual(this_stoichiometry, calculation.stoichiometry)

    def test_mul(self):
        title = "A"
        energy = -100.0
        stoichiometry = {"B": 1, "C": 2}
        calculation = Calculation(
            title=title, energy=energy, stoichiometry=stoichiometry
        )
        calculation.scale_stoichiometry = Mock(return_value={"B": 2, "C": 4})
        new_calculation = calculation * 2
        self.assertEqual(title, new_calculation.title)
        self.assertEqual(energy * 2.0, new_calculation.energy)
        self.assertEqual({"B": 2, "C": 4}, new_calculation.stoichiometry)
        calculation.scale_stoichiometry.assert_called_with(2)

    def test_truediv(self):
        title = "A"
        energy = -100.0
        stoichiometry = {"B": 4, "C": 2}
        calculation = Calculation(
            title=title, energy=energy, stoichiometry=stoichiometry
        )
        with patch("vasppy.calculation.Calculation.__mul__") as mock_mul:
            calculation / 2
            mock_mul.assert_called_with(0.5)

    def test_scale_stoichiometry(self):
        title = "A"
        energy = -100.0
        stoichiometry = {"B": 1, "C": 2}
        calculation = Calculation(
            title=title, energy=energy, stoichiometry=stoichiometry
        )
        self.assertEqual(calculation.scale_stoichiometry(2), {"B": 2, "C": 4})


class CalculationSupportFunctionsTestCase(unittest.TestCase):
    def test_delta_E(self):
        titles = ["A", "B", "C"]
        energies = [-50.5, -23.2, -10.1]
        stoichiometries = [{"B": 1, "C": 2}, {"B": 1, "C": 1}, {"C": 1}]
        calculations = [
            Calculation(title=t, energy=e, stoichiometry=s)
            for t, e, s in zip(titles, energies, stoichiometries)
        ]
        self.assertAlmostEqual(
            delta_E(reactants=[calculations[0]], products=calculations[1:3]), +17.2
        )

    @patch("vasppy.calculation.delta_stoichiometry")
    def test_delta_E_raises_value_error_if_not_balanced(self, mock_delta_stoichiometry):
        titles = ["A", "B", "C"]
        energies = [-50.5, -23.2, -10.1]
        stoichiometries = [{"B": 1, "C": 2}, {"B": 1, "C": 1}, {"C": 2}]
        calculations = [
            Calculation(title=t, energy=e, stoichiometry=s)
            for t, e, s in zip(titles, energies, stoichiometries)
        ]
        with self.assertRaises(ValueError):
            delta_E(reactants=[calculations[0]], products=calculations[1:3])
        mock_delta_stoichiometry.assert_called_with(
            [calculations[0]], calculations[1:3]
        )

    def test_delta_stoichiometry(self):
        titles = ["A", "B", "C"]
        energies = [-50.5, -23.2, -10.1]
        stoichiometries = [{"B": 1, "C": 2}, {"D": 1, "C": 1}, {"C": 1}]
        calculations = [
            Calculation(title=t, energy=e, stoichiometry=s)
            for t, e, s in zip(titles, energies, stoichiometries)
        ]
        self.assertEqual(
            delta_stoichiometry(
                reactants=[calculations[0]], products=calculations[1:3]
            ),
            {"B": -1, "D": 1},
        )

    def test_energy_string_to_float(self):
        test_strings = {"-1.2345 eV": -1.2345, "0.2341 eV": 0.2341}
        for k, v in test_strings.items():
            self.assertEqual(energy_string_to_float(k), v)

    @patch("vasppy.calculation.energy_string_to_float")
    @patch("vasppy.calculation.Calculation")
    def test_import_calculation_from_file(
        self, mock_Calculation, mock_energy_converter
    ):
        mock_energy_converter.side_effect = [-0.2414]
        example_yaml = """
        title: this_calculation
        stoichiometry:
            - A: 2
            - B: 4
        energy: -0.2414 eV
        """
        with patch("builtins.open", mock_open(read_data=example_yaml), create=True):
            import_calculations_from_file("example_file")
            mock_Calculation.assert_called_with(
                energy=-0.2414,
                stoichiometry=Counter({"B": 4, "A": 2}),
                title="this_calculation",
            )
            mock_energy_converter.assert_called_with("-0.2414 eV")

    def test_import_calculation_from_file_raises_ValueError_if_stoichiometry_is_not_set(
        self,
    ):
        example_yaml = """
        title: this_calculation
        energy: -0.2414 eV
        """
        with patch("builtins.open", mock_open(read_data=example_yaml), create=True):
            with self.assertRaises(ValueError):
                import_calculations_from_file("example_file")

    def test_import_calculation_from_file_skips_incomplete_files(self):
        example_yaml = """
        title: this_calculation
        energy: -0.2414 eV
        """
        with patch("builtins.open", mock_open(read_data=example_yaml), create=True):
            calcs = import_calculations_from_file(
                "example_file", skip_incomplete_records=True
            )
            self.assertEqual(calcs, {})


    def test_import_calculation_from_file_raises_ValueError_if_calculations_have_repeated_titles(
        self
    ):
        example_yaml = textwrap.dedent("""
        --- 
        title: this_calculation
        stoichiometry:
            - A: 2
            - B: 4
        energy: -0.2414 eV
        
        ---
        title: this_calculation
        stoichiometry:
            - A: 2
            - B: 4
        energy: -0.2414 eV
        """)
        with patch("builtins.open", mock_open(read_data=example_yaml), create=True):
            with self.assertRaises(ValueError):
                import_calculations_from_file("example_file")

                
if __name__ == "__main__":
    unittest.main()
