#include <cmath>
#include <concepts>
#include <limits>
#include <string_view>
#include <type_traits>

#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

import mean_field;
import test_helpers;

namespace {
    namespace eos     = mean_field::eos;
    namespace field   = mean_field::field;
    namespace surface = mean_field::surface;

    struct Entropy final : eos::ThermodynamicQuantity {
        static constexpr std::string_view identifier = "entropy";
    };

    struct ElectronFraction final : eos::ThermodynamicQuantity {
        static constexpr std::string_view identifier = "electron_fraction";
    };

    struct EntropyField final {
        static constexpr std::string_view name = "entropy";
    };

    struct ElectronFractionField final {
        static constexpr std::string_view name = "electron_fraction";
    };

    using SpecificEnthalpyFromPressureEntropyAndElectronFraction =
        eos::Relation<eos::quantity::SpecificEnthalpy, eos::quantity::Pressure, Entropy, ElectronFraction>;

    class GeneralStellarMatterEquationOfState final {
    public:
        using Relations = eos::RelationCatalog<SpecificEnthalpyFromPressureEntropyAndElectronFraction>;

        [[nodiscard]] constexpr eos::SpecificEnthalpyValue evaluate(
            SpecificEnthalpyFromPressureEntropyAndElectronFraction,
            const eos::PressureValue pressure,
            const eos::QuantityValue<Entropy> entropy,
            const eos::QuantityValue<ElectronFraction> electronFraction
        ) const noexcept {
            return eos::SpecificEnthalpyValue{
                2.0 * pressure.value() + 3.0 * entropy.value() + 5.0 * electronFraction.value()
            };
        }

        [[nodiscard]] constexpr eos::PartialDerivative<
            eos::quantity::SpecificEnthalpy,
            Entropy>
        partialDerivative(
            SpecificEnthalpyFromPressureEntropyAndElectronFraction,
            eos::WithRespectTo<Entropy>,
            eos::PressureValue,
            eos::QuantityValue<Entropy>,
            eos::QuantityValue<ElectronFraction>
        ) const noexcept {
            return eos::PartialDerivative<eos::quantity::SpecificEnthalpy, Entropy>{3.0};
        }

        [[nodiscard]] constexpr eos::PartialDerivative<
            eos::quantity::SpecificEnthalpy,
            ElectronFraction>
        partialDerivative(
            SpecificEnthalpyFromPressureEntropyAndElectronFraction,
            eos::WithRespectTo<ElectronFraction>,
            eos::PressureValue,
            eos::QuantityValue<Entropy>,
            eos::QuantityValue<ElectronFraction>
        ) const noexcept {
            return eos::PartialDerivative<eos::quantity::SpecificEnthalpy, ElectronFraction>{5.0};
        }
    };

    class GeneralEquationOfStateWithoutElectronFractionPartial final {
    public:
        using Relations = eos::RelationCatalog<SpecificEnthalpyFromPressureEntropyAndElectronFraction>;

        [[nodiscard]] constexpr eos::SpecificEnthalpyValue evaluate(
            SpecificEnthalpyFromPressureEntropyAndElectronFraction,
            const eos::PressureValue pressure,
            const eos::QuantityValue<Entropy> entropy,
            const eos::QuantityValue<ElectronFraction> electronFraction
        ) const noexcept {
            return eos::SpecificEnthalpyValue{pressure.value() + entropy.value() + electronFraction.value()};
        }

        [[nodiscard]] constexpr eos::PartialDerivative<
            eos::quantity::SpecificEnthalpy,
            Entropy>
        partialDerivative(
            SpecificEnthalpyFromPressureEntropyAndElectronFraction,
            eos::WithRespectTo<Entropy>,
            eos::PressureValue,
            eos::QuantityValue<Entropy>,
            eos::QuantityValue<ElectronFraction>
        ) const noexcept {
            return eos::PartialDerivative<eos::quantity::SpecificEnthalpy, Entropy>{1.0};
        }
    };

    using SpecificEnthalpyFromPressureAndEntropy =
        eos::Relation<eos::quantity::SpecificEnthalpy, eos::quantity::Pressure, Entropy>;

    class AmbiguousSurfaceEquationOfState final {
    public:
        using Relations =
            eos::RelationCatalog<eos::SpecificEnthalpyFromPressure, SpecificEnthalpyFromPressureAndEntropy>;

        [[nodiscard]] constexpr eos::SpecificEnthalpyValue evaluate(
            eos::SpecificEnthalpyFromPressure,
            const eos::PressureValue pressure
        ) const noexcept {
            return eos::SpecificEnthalpyValue{pressure.value()};
        }

        [[nodiscard]] constexpr eos::SpecificEnthalpyValue evaluate(
            SpecificEnthalpyFromPressureAndEntropy,
            const eos::PressureValue pressure,
            const eos::QuantityValue<Entropy> entropy
        ) const noexcept {
            return eos::SpecificEnthalpyValue{pressure.value() + entropy.value()};
        }

        [[nodiscard]] constexpr eos::PartialDerivative<
            eos::quantity::SpecificEnthalpy,
            Entropy>
        partialDerivative(
            SpecificEnthalpyFromPressureAndEntropy,
            eos::WithRespectTo<Entropy>,
            eos::PressureValue,
            eos::QuantityValue<Entropy>
        ) const noexcept {
            return eos::PartialDerivative<eos::quantity::SpecificEnthalpy, Entropy>{1.0};
        }
    };

    class DensityOnlyEquationOfState final {
    public:
        using Relations = eos::RelationCatalog<eos::PressureFromDensity>;

        [[nodiscard]] constexpr eos::PressureValue evaluate(
            eos::PressureFromDensity,
            const eos::DensityValue density
        ) const noexcept {
            return eos::PressureValue{density.value()};
        }
    };

    using GeneralSurfaceFormulation = surface::SurfaceConstraintFormulation<
        eos::quantity::SpecificEnthalpy,
        field::Enthalpy,
        surface::SurfaceStateBindings<
            surface::SurfaceStateBinding<eos::quantity::SpecificEnthalpy, field::Enthalpy>,
            surface::SurfaceStateBinding<Entropy, EntropyField>,
            surface::SurfaceStateBinding<ElectronFraction, ElectronFractionField>>>;

    struct PolytropicSurfaceState final {
        double specificEnthalpy;

        [[nodiscard]] eos::SpecificEnthalpyValue value(eos::quantity::SpecificEnthalpy) const noexcept {
            return eos::SpecificEnthalpyValue{specificEnthalpy};
        }
    };

    struct GeneralSurfaceState final {
        double specificEnthalpy;
        double entropy;
        double electronFraction;

        [[nodiscard]] eos::SpecificEnthalpyValue value(eos::quantity::SpecificEnthalpy) const noexcept {
            return eos::SpecificEnthalpyValue{specificEnthalpy};
        }

        [[nodiscard]] eos::QuantityValue<Entropy> value(Entropy) const noexcept {
            return eos::QuantityValue<Entropy>{entropy};
        }

        [[nodiscard]] eos::QuantityValue<ElectronFraction> value(ElectronFraction) const noexcept {
            return eos::QuantityValue<ElectronFraction>{electronFraction};
        }
    };

    template <typename Candidate>
    concept HasTargetEnthalpy = requires(const Candidate &candidate) { candidate.targetEnthalpy; };
} // namespace

TEST_CASE(
    "Constant Pressure Surface Prescribes Only A Pressure Quantity",
    tags::surface_prescription_type_contract
) {
    STATIC_CHECK(std::same_as<surface::ConstantPressureSurface::PhysicalQuantity, eos::quantity::Pressure>);
    STATIC_CHECK(std::constructible_from<surface::ConstantPressureSurface, eos::PressureValue>);
    STATIC_CHECK_FALSE(std::constructible_from<surface::ConstantPressureSurface, eos::SpecificEnthalpyValue>);
    STATIC_CHECK_FALSE(std::constructible_from<surface::ConstantPressureSurface, double>);
    STATIC_CHECK(std::same_as<surface::Isobaric, surface::ConstantPressureSurface>);
    STATIC_CHECK(std::is_trivially_copyable_v<surface::ConstantPressureSurface>);
    STATIC_CHECK(std::is_trivially_copyable_v<surface::PressureSurfaceDescriptor>);
    STATIC_CHECK(std::is_trivially_copyable_v<surface::RuntimeSurfaceConstraintDependencies>);

    const surface::ConstantPressureSurface pressureSurface{eos::PressureValue{0.03125}};

    CHECK(pressureSurface.targetPressure() == eos::PressureValue{0.03125});
    CHECK(pressureSurface.descriptor().targetPressure == 0.03125);

    CHECK_THROWS_AS(surface::ConstantPressureSurface{eos::PressureValue{-0.1}}, std::invalid_argument);
    CHECK_THROWS_AS(
        surface::ConstantPressureSurface{eos::PressureValue{std::numeric_limits<double>::infinity()}},
        std::invalid_argument
    );
}

TEST_CASE(
    "Polytropic EOS Resolves Constant Surface Pressure Through Its Enthalpy Relation",
    tags::surface_constraint_compilation
) {
    using Formulation = surface::BarotropicSurfaceFormulation;

    STATIC_CHECK(surface::PressureSurfaceCompilable<Formulation, eos::Polytrope>);
    STATIC_CHECK_FALSE(surface::PressureSurfaceCompilable<Formulation, DensityOnlyEquationOfState>);

    const eos::Polytrope equationOfState(3.0, 0.25);
    const surface::ConstantPressureSurface pressureSurface{eos::PressureValue{0.03125}};
    const auto constraint = surface::compilePressureSurfaceConstraint<Formulation>(pressureSurface, equationOfState);

    using Constraint      = std::remove_cvref_t<decltype(constraint)>;
    using Dependencies    = Constraint::SurfaceDependencies;

    STATIC_CHECK(std::is_trivially_copyable_v<Constraint>);
    STATIC_CHECK(std::same_as<Constraint::Relation, eos::SpecificEnthalpyFromPressure>);
    STATIC_CHECK(std::same_as<Dependencies::RowField, field::Enthalpy>);
    STATIC_CHECK(std::same_as<Dependencies::StateFieldTypes, field::TypeList<field::Enthalpy>>);
    STATIC_CHECK_FALSE(HasTargetEnthalpy<Constraint>);

    const double requiredSpecificEnthalpy =
        eos::evaluate<eos::quantity::SpecificEnthalpy>(equationOfState, pressureSurface.targetPressure()).value();
    const PolytropicSurfaceState state{requiredSpecificEnthalpy};
    const PolytropicSurfaceState variation{-0.19};

    CHECK(constraint.targetPressure() == eos::PressureValue{0.03125});
    CHECK(constraint.residual(state) == 0.0);
    CHECK(constraint.jacobianAction(state, variation) == -0.19);

    const auto runtimeDependencies = constraint.runtimeDependencies();
    REQUIRE(runtimeDependencies.stateFields.size() == 1);
    CHECK(runtimeDependencies.residualRowField == surface::surfaceFieldId<field::Enthalpy>);
    CHECK(runtimeDependencies.stateFields[0] == surface::surfaceFieldId<field::Enthalpy>);
}

TEST_CASE(
    "General EOS Resolves Constant Surface Pressure With Local Composition",
    tags::surface_constraint_compilation
) {
    STATIC_CHECK(surface::PressureSurfaceCompilable<GeneralSurfaceFormulation, GeneralStellarMatterEquationOfState>);
    STATIC_CHECK_FALSE(
        surface::PressureSurfaceCompilable<surface::BarotropicSurfaceFormulation, GeneralStellarMatterEquationOfState>
    );
    STATIC_CHECK_FALSE(
        surface::PressureSurfaceCompilable<
            GeneralSurfaceFormulation, GeneralEquationOfStateWithoutElectronFractionPartial>
    );
    STATIC_CHECK_FALSE(surface::PressureSurfaceCompilable<GeneralSurfaceFormulation, AmbiguousSurfaceEquationOfState>);

    const GeneralStellarMatterEquationOfState equationOfState;
    const surface::ConstantPressureSurface pressureSurface{eos::PressureValue{0.4}};
    const auto constraint =
        surface::compilePressureSurfaceConstraint<GeneralSurfaceFormulation>(pressureSurface, equationOfState);

    using Constraint   = std::remove_cvref_t<decltype(constraint)>;
    using Dependencies = Constraint::SurfaceDependencies;

    STATIC_CHECK(std::same_as<Constraint::Relation, SpecificEnthalpyFromPressureEntropyAndElectronFraction>);
    STATIC_CHECK(
        std::same_as<
            Dependencies::StateFieldTypes, field::TypeList<field::Enthalpy, EntropyField, ElectronFractionField>>
    );

    constexpr GeneralSurfaceState firstSurface{
        .specificEnthalpy = 2.0 * 0.4 + 3.0 * 0.2 + 5.0 * 0.1, .entropy = 0.2, .electronFraction = 0.1
    };
    constexpr GeneralSurfaceState secondSurface{
        .specificEnthalpy = 2.0 * 0.4 + 3.0 * 0.3 + 5.0 * 0.1, .entropy = 0.3, .electronFraction = 0.1
    };

    CHECK(firstSurface.specificEnthalpy != secondSurface.specificEnthalpy);
    CHECK(constraint.residual(firstSurface) == 0.0);
    CHECK(constraint.residual(secondSurface) == 0.0);

    const auto runtimeDependencies = constraint.runtimeDependencies();
    REQUIRE(runtimeDependencies.stateFields.size() == 3);
    CHECK(runtimeDependencies.stateFields[0] == surface::surfaceFieldId<field::Enthalpy>);
    CHECK(runtimeDependencies.stateFields[1] == surface::surfaceFieldId<EntropyField>);
    CHECK(runtimeDependencies.stateFields[2] == surface::surfaceFieldId<ElectronFractionField>);
}

TEST_CASE(
    "General EOS Pressure Surface Jacobian Includes Every Local State Dependency",
    tags::surface_constraint_jacobian
) {
    const GeneralStellarMatterEquationOfState equationOfState;
    const surface::ConstantPressureSurface pressureSurface{eos::PressureValue{0.4}};
    const auto constraint =
        surface::compilePressureSurfaceConstraint<GeneralSurfaceFormulation>(pressureSurface, equationOfState);

    constexpr GeneralSurfaceState state{.specificEnthalpy = 1.7, .entropy = 0.2, .electronFraction = 0.1};
    constexpr GeneralSurfaceState variation{.specificEnthalpy = 0.7, .entropy = -0.2, .electronFraction = 0.05};
    constexpr double step = 1.0e-7;

    const GeneralSurfaceState forward{
        .specificEnthalpy = state.specificEnthalpy + step * variation.specificEnthalpy,
        .entropy          = state.entropy + step * variation.entropy,
        .electronFraction = state.electronFraction + step * variation.electronFraction
    };
    const GeneralSurfaceState backward{
        .specificEnthalpy = state.specificEnthalpy - step * variation.specificEnthalpy,
        .entropy          = state.entropy - step * variation.entropy,
        .electronFraction = state.electronFraction - step * variation.electronFraction
    };

    const double finiteDifference = (constraint.residual(forward) - constraint.residual(backward)) / (2.0 * step);
    const double jacobianAction   = constraint.jacobianAction(state, variation);

    CHECK(jacobianAction == variation.specificEnthalpy - 3.0 * variation.entropy - 5.0 * variation.electronFraction);
    CHECK_THAT(finiteDifference, Catch::Matchers::WithinAbs(jacobianAction, 2.0e-9));
}