#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <limits>
#include <type_traits>
#include <utility>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

import mean_field;
import mean_field_extension_example.ideal_gas_radiation;

namespace {
    namespace dimensions = mean_field::dimensions;
    namespace eos        = mean_field::eos;
    namespace example    = mean_field::extension_example;

    [[nodiscard]] example::IdealGasRadiation makeSimpleEquationOfState() {
        /* R = k_B / (mu m_u) = 12 / (2 * 3) = 2. */
        return example::IdealGasRadiation({
            .meanMolecularWeight = 2.0,
            .boltzmannConstant = 12.0,
            .atomicMassUnit = 3.0,
            .radiationConstant = 9.0
        });
    }

    template <typename Function>
    [[nodiscard]] double centeredDifference(
        Function function,
        const double point
    ) {
        const double step = std::cbrt(std::numeric_limits<double>::epsilon()) *
                            std::max(1.0, std::abs(point));
        return (function(point + step) - function(point - step)) / (2.0 * step);
    }

    template <typename EquationOfState>
    concept CanEvaluatePressureWithReversedInputs = requires(
        const EquationOfState &equationOfState,
        const dimensions::TemperatureValue temperature,
        const dimensions::DensityValue density
    ) {
        eos::evaluate<dimensions::quantity::Pressure>(equationOfState, temperature, density);
    };
} // namespace

TEST_CASE("The extension satisfies the EOS protocol at compile time", "[extension-example][eos][type]") {
    using EquationOfState = example::IdealGasRadiation;

    STATIC_CHECK(mean_field::models::SelfDescribingModelSpecification<EquationOfState>);
    STATIC_CHECK(eos::EquationOfStateModel<EquationOfState>);
    STATIC_CHECK(eos::SupportsRelation<EquationOfState, example::PressureFromDensityAndTemperature>);
    STATIC_CHECK(eos::SupportsRelation<EquationOfState, example::SpecificInternalEnergyFromDensityAndTemperature>);
    STATIC_CHECK(eos::SupportsRelation<EquationOfState, example::SpecificEnthalpyFromDensityAndTemperature>);
    STATIC_CHECK_FALSE(eos::BarotropicClosureEquationOfState<EquationOfState>);
    STATIC_CHECK_FALSE(CanEvaluatePressureWithReversedInputs<EquationOfState>);

    using PressureResult = decltype(eos::evaluate<dimensions::quantity::Pressure>(
        std::declval<const EquationOfState &>(),
        dimensions::DensityValue{1.0},
        dimensions::TemperatureValue{1.0}
    ));
    STATIC_CHECK(std::same_as<PressureResult, dimensions::PressureValue>);
}

TEST_CASE("Gas and radiation terms reproduce the defining thermodynamics", "[extension-example][eos][physics]") {
    const auto equationOfState = makeSimpleEquationOfState();
    const dimensions::DensityValue density{4.0};
    const dimensions::TemperatureValue temperature{2.0};

    const auto pressureContributions = equationOfState.pressureContributions(density, temperature);
    const auto pressure = eos::evaluate<dimensions::quantity::Pressure>(
        equationOfState,
        density,
        temperature
    );
    const auto internalEnergy = eos::evaluate<dimensions::quantity::SpecificInternalEnergy>(
        equationOfState,
        density,
        temperature
    );
    const auto enthalpy = eos::evaluate<dimensions::quantity::SpecificEnthalpy>(
        equationOfState,
        density,
        temperature
    );

    CHECK(equationOfState.specificGasConstant() == Catch::Approx(2.0));
    CHECK(pressureContributions.gas.value() == Catch::Approx(16.0));
    CHECK(pressureContributions.radiation.value() == Catch::Approx(48.0));
    CHECK(pressure.value() == Catch::Approx(64.0));
    CHECK(internalEnergy.value() == Catch::Approx(42.0));
    CHECK(enthalpy.value() == Catch::Approx(58.0));

    /* This is the thermodynamic identity h = u + P/rho. */
    CHECK(enthalpy.value() == Catch::Approx(internalEnergy.value() + pressure.value() / density.value()));
}

TEST_CASE("The gas and photon terms have their expected scaling laws", "[extension-example][eos][physics]") {
    const auto equationOfState = makeSimpleEquationOfState();
    const dimensions::DensityValue density{3.5};
    const dimensions::TemperatureValue temperature{1.25};

    const auto baseline = equationOfState.pressureContributions(density, temperature);
    const auto doubledDensity = equationOfState.pressureContributions(
        dimensions::DensityValue{2.0 * density.value()},
        temperature
    );
    const auto doubledTemperature = equationOfState.pressureContributions(
        density,
        dimensions::TemperatureValue{2.0 * temperature.value()}
    );

    CHECK(doubledDensity.gas.value() == Catch::Approx(2.0 * baseline.gas.value()));
    CHECK(doubledDensity.radiation.value() == Catch::Approx(baseline.radiation.value()));
    CHECK(doubledTemperature.gas.value() == Catch::Approx(2.0 * baseline.gas.value()));
    CHECK(doubledTemperature.radiation.value() == Catch::Approx(16.0 * baseline.radiation.value()));

    const double crossoverTemperature = std::cbrt(
        3.0 * density.value() * equationOfState.specificGasConstant() /
        equationOfState.parameters().radiationConstant
    );
    const auto crossover = equationOfState.pressureContributions(
        density,
        dimensions::TemperatureValue{crossoverTemperature}
    );
    CHECK(crossover.gas.value() == Catch::Approx(crossover.radiation.value()).epsilon(2.0e-14));
}

TEST_CASE("All declared Jacobian entries match centered numerical derivatives",
          "[extension-example][eos][derivative][numerical]") {
    const auto equationOfState = example::IdealGasRadiation({
        .meanMolecularWeight = 1.25,
        .boltzmannConstant = 2.75,
        .atomicMassUnit = 0.8,
        .radiationConstant = 0.35
    });

    struct State final {
        double density;
        double temperature;
    };
    const std::array states{
        State{.density = 0.4, .temperature = 0.7},
        State{.density = 2.0, .temperature = 1.5},
        State{.density = 11.0, .temperature = 3.0}
    };

    for (const State state : states) {
        const dimensions::DensityValue density{state.density};
        const dimensions::TemperatureValue temperature{state.temperature};

        const auto pressureDensity = eos::partialDerivative<
            dimensions::quantity::Pressure,
            dimensions::quantity::Density>(equationOfState, density, temperature);
        const auto pressureTemperature = eos::partialDerivative<
            dimensions::quantity::Pressure,
            dimensions::quantity::Temperature>(equationOfState, density, temperature);
        const auto energyDensity = eos::partialDerivative<
            dimensions::quantity::SpecificInternalEnergy,
            dimensions::quantity::Density>(equationOfState, density, temperature);
        const auto energyTemperature = eos::partialDerivative<
            dimensions::quantity::SpecificInternalEnergy,
            dimensions::quantity::Temperature>(equationOfState, density, temperature);
        const auto enthalpyDensity = eos::partialDerivative<
            dimensions::quantity::SpecificEnthalpy,
            dimensions::quantity::Density>(equationOfState, density, temperature);
        const auto enthalpyTemperature = eos::partialDerivative<
            dimensions::quantity::SpecificEnthalpy,
            dimensions::quantity::Temperature>(equationOfState, density, temperature);

        const double numericalPressureDensity = centeredDifference(
            [&](const double rho) {
                return eos::evaluate<dimensions::quantity::Pressure>(
                    equationOfState,
                    dimensions::DensityValue{rho},
                    temperature
                ).value();
            },
            state.density
        );
        const double numericalPressureTemperature = centeredDifference(
            [&](const double T) {
                return eos::evaluate<dimensions::quantity::Pressure>(
                    equationOfState,
                    density,
                    dimensions::TemperatureValue{T}
                ).value();
            },
            state.temperature
        );
        const double numericalEnergyDensity = centeredDifference(
            [&](const double rho) {
                return eos::evaluate<dimensions::quantity::SpecificInternalEnergy>(
                    equationOfState,
                    dimensions::DensityValue{rho},
                    temperature
                ).value();
            },
            state.density
        );
        const double numericalEnergyTemperature = centeredDifference(
            [&](const double T) {
                return eos::evaluate<dimensions::quantity::SpecificInternalEnergy>(
                    equationOfState,
                    density,
                    dimensions::TemperatureValue{T}
                ).value();
            },
            state.temperature
        );
        const double numericalEnthalpyDensity = centeredDifference(
            [&](const double rho) {
                return eos::evaluate<dimensions::quantity::SpecificEnthalpy>(
                    equationOfState,
                    dimensions::DensityValue{rho},
                    temperature
                ).value();
            },
            state.density
        );
        const double numericalEnthalpyTemperature = centeredDifference(
            [&](const double T) {
                return eos::evaluate<dimensions::quantity::SpecificEnthalpy>(
                    equationOfState,
                    density,
                    dimensions::TemperatureValue{T}
                ).value();
            },
            state.temperature
        );

        constexpr double tolerance = 3.0e-9;
        CHECK(pressureDensity.value() == Catch::Approx(numericalPressureDensity).epsilon(tolerance));
        CHECK(pressureTemperature.value() == Catch::Approx(numericalPressureTemperature).epsilon(tolerance));
        CHECK(energyDensity.value() == Catch::Approx(numericalEnergyDensity).epsilon(tolerance));
        CHECK(energyTemperature.value() == Catch::Approx(numericalEnergyTemperature).epsilon(tolerance));
        CHECK(enthalpyDensity.value() == Catch::Approx(numericalEnthalpyDensity).epsilon(tolerance));
        CHECK(enthalpyTemperature.value() == Catch::Approx(numericalEnthalpyTemperature).epsilon(tolerance));
    }
}

TEST_CASE("The physical domain is checked at the EOS boundary", "[extension-example][eos][domain]") {
    const auto equationOfState = makeSimpleEquationOfState();
    const double nan = std::numeric_limits<double>::quiet_NaN();

    CHECK_THROWS_AS(
        example::IdealGasRadiation({
            .meanMolecularWeight = 0.0,
            .boltzmannConstant = 1.0,
            .atomicMassUnit = 1.0,
            .radiationConstant = 1.0
        }),
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        example::IdealGasRadiation({
            .meanMolecularWeight = 1.0,
            .boltzmannConstant = 1.0,
            .atomicMassUnit = 1.0,
            .radiationConstant = -1.0
        }),
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        eos::evaluate<dimensions::quantity::Pressure>(
            equationOfState,
            dimensions::DensityValue{0.0},
            dimensions::TemperatureValue{1.0}
        ),
        eos::EvaluationError
    );
    CHECK_THROWS_AS(
        eos::evaluate<dimensions::quantity::Pressure>(
            equationOfState,
            dimensions::DensityValue{1.0},
            dimensions::TemperatureValue{-1.0}
        ),
        eos::EvaluationError
    );
    CHECK_THROWS_AS(
        eos::evaluate<dimensions::quantity::Pressure>(
            equationOfState,
            dimensions::DensityValue{nan},
            dimensions::TemperatureValue{1.0}
        ),
        eos::EvaluationError
    );
}