from __future__ import annotations

import unittest
from unittest.mock import Mock, patch

from app.simulation.components.amesim.media.mediums import (
    AmesimHeliumPengRobinsonMedium,
)
from app.simulation.property_cache import property_cache_run


class PropertyPerformanceInstrumentationTests(unittest.TestCase):
    def setUp(self) -> None:
        self.medium = AmesimHeliumPengRobinsonMedium()

    def test_temperature_iterations_are_recorded_only_on_cache_miss(self) -> None:
        pressure = 15.3e6
        temperature = 293.15
        enthalpy = self.medium.specific_enthalpy_at_pressure(
            pressure,
            temperature,
        )

        with patch(
            "app.simulation.components.amesim.media.mediums."
            "record_property_iterations"
        ) as record_iterations:
            with property_cache_run():
                first = self.medium.temperature_from_pressure_enthalpy(
                    pressure,
                    enthalpy,
                )
                second = self.medium.temperature_from_pressure_enthalpy(
                    pressure,
                    enthalpy,
                )

        self.assertEqual(second, first)
        record_iterations.assert_called_once()
        operation, iterations, converged = record_iterations.call_args.args
        self.assertEqual(operation, "temperature_from_pressure_enthalpy")
        self.assertEqual(iterations, 5)
        self.assertTrue(converged)

    def test_state_recovery_iterations_are_recorded_only_on_cache_miss(self) -> None:
        pressure = 15.3e6
        temperature = 293.15
        volume = 0.057
        mass = self.medium.density(pressure, temperature) * volume
        internal_energy = mass * self.medium.specific_internal_energy_at_pressure(
            pressure,
            temperature,
        )

        with patch(
            "app.simulation.components.amesim.media.mediums."
            "record_property_iterations"
        ) as record_iterations:
            with property_cache_run():
                first = self.medium.properties_from_mU(
                    mass,
                    internal_energy,
                    volume,
                )
                second = self.medium.properties_from_mU(
                    mass,
                    internal_energy,
                    volume,
                )

        self.assertIs(second, first)
        record_iterations.assert_called_once()
        operation, iterations, converged = record_iterations.call_args.args
        self.assertEqual(operation, "properties_from_mU")
        self.assertEqual(iterations, 5)
        self.assertTrue(converged)

    def test_temperature_iteration_limit_is_reported_as_not_converged(self) -> None:
        pressure = 123_456.0
        base_temperature = 300.0
        enthalpy = self.medium.specific_enthalpy(base_temperature)
        fake_fluid = Mock()
        fake_fluid.residual_specific_enthalpy.side_effect = (
            lambda _pressure, temperature: (
                self.medium.cp_ref * 10.0
                if temperature >= base_temperature
                else -self.medium.cp_ref * 10.0
            )
        )

        with (
            patch.object(AmesimHeliumPengRobinsonMedium, "fluid", fake_fluid),
            patch(
                "app.simulation.components.amesim.media.mediums."
                "record_property_iterations"
            ) as record_iterations,
        ):
            self.medium.temperature_from_pressure_enthalpy(
                pressure,
                enthalpy,
            )

        record_iterations.assert_called_once_with(
            "temperature_from_pressure_enthalpy",
            16,
            False,
        )

    def test_state_recovery_iteration_limit_is_reported_as_not_converged(self) -> None:
        base_temperature = 300.0
        target_internal_energy = self.medium.specific_internal_energy(
            base_temperature
        )
        fake_fluid = Mock()
        fake_fluid.residual_specific_internal_energy_at_density.side_effect = (
            lambda temperature, _density: (
                self.medium.cv * 10.0
                if temperature >= base_temperature
                else -self.medium.cv * 10.0
            )
        )
        fake_fluid.pressure_from_density.return_value = 101_325.0
        fake_fluid.residual_specific_enthalpy.return_value = 0.0

        with (
            patch.object(AmesimHeliumPengRobinsonMedium, "fluid", fake_fluid),
            patch(
                "app.simulation.components.amesim.media.mediums."
                "record_property_iterations"
            ) as record_iterations,
        ):
            self.medium.properties_from_mU(
                1.0,
                target_internal_energy,
                1.0,
            )

        record_iterations.assert_called_once_with(
            "properties_from_mU",
            16,
            False,
        )


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