#include <cmath>
#include <limits>
#include <memory>
#include <string>
#include <type_traits>
#include <vector>

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

#include "serif/model/description.hpp"
#include "serif/model/eos_reference.hpp"
#include "serif/tests/equations/test_equations.hpp"
#include "serif/equations/registry.hpp"
#include "serif/geometry/surface/isobaric.hpp"

namespace
{
    namespace model = serif::model;
    namespace dimensions = serif::dimensions;

    struct SymbolicEOS
    {
        double parameter;
    };

    struct MoveOnlyEOS
    {
        std::unique_ptr<double> parameter;
    };

    struct MoveOnlyConstraint
    {
        using ScalarDescription = model::ScalarConstraint<dimensions::Mass, dimensions::SpecificEnthalpy, dimensions::Mass,
                                                          "move.coordinate", "C", "move.residual", "R">;
        using ModelDefinition = model::SpecificationDefinition<MoveOnlyConstraint, "MoveOnlyConstraint", model::SpecificationRole::invariant,
                                                              model::GeneratedStateKind::multiplier, ScalarDescription>;
        std::unique_ptr<double> mass;
        /** @brief Expose a dimensional target while preserving unique ownership. */
        dimensions::MassValue target() const
        {
            return dimensions::MassValue{*mass};
        }
    };

    struct MoveOnlyPhysicalEOS
    {
        using Relation = serif::eos::relations::DensityFromSpecificEnthalpy;
        using Relations = serif::eos::relations::RelationCatalog<Relation>;
        std::unique_ptr<double> coefficient;

        /** @brief Supply an independent linear closure with a uniquely owned material parameter. */
        dimensions::DensityValue evaluate(Relation, dimensions::SpecificEnthalpyValue enthalpy) const noexcept
        {
            return dimensions::DensityValue{*coefficient * enthalpy.value()};
        }

        /** @brief Supply the exact derivative while retaining the EOS relation contract. */
        dimensions::PartialDerivative<dimensions::MassDensity, dimensions::SpecificEnthalpy> partial_derivative(Relation, dimensions::WithRespectTo<dimensions::SpecificEnthalpy>, dimensions::SpecificEnthalpyValue) const noexcept
        {
            return dimensions::PartialDerivative<dimensions::MassDensity, dimensions::SpecificEnthalpy>{*coefficient};
        }
    };

    /** @brief Preserve surface members through aggregate deduction and the substitution performed by a templated factory. */
    template <typename EOS>
    [[nodiscard]] auto model_with_surface(EOS eos, dimensions::PressureValue target)
    {
        return model::make_stellar_model(std::move(eos), model::EquilibriumEquations
        {
            .interior_equations = serif::equations::EquationRegistry{::serif::tests::equations::BarotropicClosure{serif::eos::models::Polytrope{2.0, 0.3}}},
            .shape_force = model::NoShapeForce{},
            .surface_condition = serif::geometry::surface::Isobaric{target}
        });
    }
}

TEST_CASE("Templated model factories retain the declared surface type and physical pressure", "[model][unit][surface]")
{
    const auto star = model_with_surface(SymbolicEOS{3.0}, dimensions::PressureValue{0.8});
    using Description = typename decltype(star)::Equations;
    using Surface = serif::geometry::surface::Isobaric;
    STATIC_CHECK(std::same_as<typename Description::SurfaceCondition, Surface>);
    STATIC_CHECK(std::same_as<typename Description::ShapeForce, model::NoShapeForce>);
    STATIC_CHECK(std::same_as<std::remove_cvref_t<decltype(star.equations().surface_condition)>, Surface>);
    CHECK(star.equations().surface_condition.target_pressure().value() == 0.8);
    const auto equations = model::resolve_equilibrium_equations(star);
    STATIC_CHECK(std::same_as<typename decltype(equations)::SurfaceCondition, Surface>);
    CHECK(&equations.surface_condition == &star.equations().surface_condition);
    CHECK(equations.surface_condition.target_pressure().value() == 0.8);
}

TEST_CASE("Stellar descriptions retain typed physics in canonical order", "[model][unit]")
{
    const auto star = model::make_stellar_model(SymbolicEOS{3.0}, model::EquilibriumEquations{17},
        model::FixedAngularMomentum{dimensions::AngularMomentumValue{0.0}},
        model::FixedTotalMass{dimensions::MassValue{1.5}});
    CHECK(star.equation_of_state().parameter == 3.0);
    CHECK(star.equations().interior_equations == 17);
    CHECK(star.specification<model::FixedTotalMass>().target().value() == 1.5);
    CHECK(std::tuple_size_v<decltype(star.specifications_for_role<model::SpecificationRole::invariant>())> == 2);
    const auto descriptors = star.specification_descriptors();
    REQUIRE(descriptors.size() == 2);
    CHECK(descriptors[0].name == "FixedTotalMass");
    CHECK(descriptors[1].name == "FixedAngularMomentum");

    std::vector<double> targets;
    star.visit_specifications([&](const auto& specification)
    {
        if constexpr (requires { specification.target(); })
        {
            targets.push_back(specification.target().value());
        }
        else
        {
            targets.push_back(specification.parameter);
        }
    });
    CHECK(targets == std::vector<double>{3.0, 1.5, 0.0});
}

TEST_CASE("Explicit equation definitions read one authoritative EOS and target at binding", "[model][unit]")
{
    int resolutions = 0;
    const auto definition = [&resolutions](const auto& physics)
    {
        ++resolutions;
        return model::EquilibriumEquations{physics.equation_of_state().parameter + physics.template specification<model::FixedTotalMass>().target().value()};
    };
    const auto star = model::make_stellar_model(SymbolicEOS{3.0}, definition, model::FixedTotalMass{dimensions::MassValue{2.0}});
    CHECK(resolutions == 0);
    const auto equations = model::resolve_equilibrium_equations(star);
    CHECK(resolutions == 1);
    CHECK(equations.interior_equations == 5.0);
    const auto changed = star.transform_specifications([](const auto& physics)
    {
        using Physics = std::remove_cvref_t<decltype(physics)>;
        if constexpr (std::same_as<Physics, model::FixedTotalMass>)
        {
            return Physics{dimensions::MassValue{7.0}};
        }
        else
        {
            return physics;
        }
    });
    CHECK(model::resolve_equilibrium_equations(changed).interior_equations == 10.0);
    CHECK(model::resolve_equilibrium_equations(star).interior_equations == 5.0);
}

TEST_CASE("Equation composition resolves member definitions and borrows move-only stored equations", "[model][unit]")
{
    const auto definitions = model::EquilibriumEquations
    {
        .interior_equations = [](const auto& physics)
        {
            return std::make_unique<double>(physics.equation_of_state().parameter);
        },
        .shape_force = std::make_unique<int>(8)
    };
    const auto star = model::make_stellar_model(SymbolicEOS{4.0}, model::EquilibriumEquations
    {
        .interior_equations = definitions.interior_equations,
        .shape_force = std::make_unique<int>(8)
    });
    const auto equations = model::resolve_equilibrium_equations(star);
    STATIC_CHECK(std::is_reference_v<decltype(equations.shape_force)>);
    STATIC_CHECK_FALSE(std::is_reference_v<decltype(equations.interior_equations)>);
    CHECK(*equations.interior_equations == 4.0);
    CHECK(&equations.shape_force == &star.equations().shape_force);
    CHECK(*equations.shape_force == 8);
}

TEST_CASE("Move-only physical prescriptions survive construction and explicit reconstruction", "[model][unit]")
{
    const auto star = model::make_stellar_model(MoveOnlyEOS{std::make_unique<double>(3.0)}, model::EquilibriumEquations{std::make_unique<int>(5)},
        MoveOnlyConstraint{std::make_unique<double>(2.0)}, model::FixedTotalMass{dimensions::MassValue{7.0}});
    CHECK(*star.equation_of_state().parameter == 3.0);
    CHECK(star.specification<MoveOnlyConstraint>().target().value() == 2.0);
    const auto changed = star.transform_specifications([](const auto& physics)
    {
        using Physics = std::remove_cvref_t<decltype(physics)>;
        if constexpr (std::same_as<Physics, MoveOnlyEOS>)
        {
            return Physics{std::make_unique<double>(*physics.parameter)};
        }
        else if constexpr (std::same_as<Physics, MoveOnlyConstraint>)
        {
            return Physics{std::make_unique<double>(11.0)};
        }
        else
        {
            return physics;
        }
    }, model::EquilibriumEquations{std::make_unique<int>(9)});
    CHECK(changed.specification<MoveOnlyConstraint>().target().value() == 11.0);
    CHECK(star.specification<MoveOnlyConstraint>().target().value() == 2.0);
    CHECK(*model::resolve_equilibrium_equations(changed).interior_equations == 9);
}

TEST_CASE("Fixed physical targets reject invalid numbers and retain independent coordinate units", "[model][unit][validation]")
{
    for (double invalid : {0.0, -1.0, std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()})
    {
        CHECK_THROWS_AS(model::FixedTotalMass{dimensions::MassValue{invalid}}, std::invalid_argument);
        if (invalid != 0.0)
        {
            CHECK_THROWS_AS(model::FixedAngularMomentum{dimensions::AngularMomentumValue{invalid}}, std::invalid_argument);
        }
    }
    CHECK_NOTHROW(model::FixedAngularMomentum{dimensions::AngularMomentumValue{0.0}});
    CHECK_THROWS_AS(model::FixedAngularMomentum(model::FixedAngularMomentum::Parameters{.angular_momentum = dimensions::AngularMomentumValue{1.0}, .axis = {0.0, 0.0, 0.0}}), std::invalid_argument);
    for (double invalid : {std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()})
    {
        for (std::size_t component = 0; component < 3; ++component)
        {
            model::FixedAngularMomentum::Parameters invalid_axis{.angular_momentum = dimensions::AngularMomentumValue{1.0}};
            invalid_axis.axis[component] = invalid;
            CHECK_THROWS_AS(model::FixedAngularMomentum{invalid_axis}, std::invalid_argument);
            model::FixedAngularMomentum::Parameters invalid_center{.angular_momentum = dimensions::AngularMomentumValue{1.0}};
            invalid_center.center[component] = invalid;
            CHECK_THROWS_AS(model::FixedAngularMomentum{invalid_center}, std::invalid_argument);
        }
    }
    const model::FixedAngularMomentum prescribed(model::FixedAngularMomentum::Parameters{.angular_momentum = dimensions::AngularMomentumValue{2.0}, .axis = {3.0, 4.0, 0.0}, .center = {1.0, 2.0, 3.0}});
    CHECK(prescribed.axis()[0] == Catch::Approx(0.6));
    CHECK(prescribed.axis()[1] == Catch::Approx(0.8));
    CHECK(prescribed.center() == std::array<double, 3>{1.0, 2.0, 3.0});
    for (double scale : {std::numeric_limits<double>::max(), std::numeric_limits<double>::denorm_min()})
    {
        const model::FixedAngularMomentum extreme(model::FixedAngularMomentum::Parameters{.angular_momentum = dimensions::AngularMomentumValue{0.0}, .axis = {scale, scale, 0.0}});
        CHECK(extreme.axis()[0] == Catch::Approx(std::sqrt(0.5)));
        CHECK(extreme.axis()[1] == Catch::Approx(std::sqrt(0.5)));
    }
}

TEST_CASE("Typed EOS references retain owners and exact relation capabilities", "[model][unit][eos]")
{
    auto star = model::make_stellar_model(MoveOnlyPhysicalEOS{std::make_unique<double>(2.0)}, model::EquilibriumEquations{1});
    using Model = decltype(star);
    std::shared_ptr<const Model> owner = std::make_shared<Model>(std::move(star));
    std::weak_ptr<const Model> lifetime = owner;
    const auto shared = model::share_equation_of_state(owner);
    STATIC_CHECK(serif::eos::models::EOSSupportingPartialDerivative<decltype(shared), serif::eos::relations::DensityFromSpecificEnthalpy, dimensions::SpecificEnthalpy>);
    CHECK(&shared.get() == &owner->equation_of_state());
    owner.reset();
    CHECK_FALSE(lifetime.expired());
    CHECK(serif::eos::evaluate<dimensions::MassDensity>(shared, dimensions::SpecificEnthalpyValue{3.0}).value() == 6.0);
    CHECK(serif::eos::partial_derivative<dimensions::MassDensity, dimensions::SpecificEnthalpy>(shared, dimensions::SpecificEnthalpyValue{3.0}).value() == 2.0);
    STATIC_CHECK(noexcept(shared.evaluate(MoveOnlyPhysicalEOS::Relation{}, dimensions::SpecificEnthalpyValue{1.0})));
    CHECK_THROWS_AS(model::EosReference<MoveOnlyPhysicalEOS>(std::shared_ptr<const MoveOnlyPhysicalEOS>{}), std::invalid_argument);
    CHECK_THROWS_AS(model::share_equation_of_state(std::shared_ptr<const Model>{}), std::invalid_argument);
}

TEST_CASE("Programmer supplied physical equation recipes share the model EOS", "[model][unit][eos]")
{
    const auto star = model::make_stellar_model(MoveOnlyPhysicalEOS{std::make_unique<double>(4.0)}, model::EquilibriumEquations
    {
        .interior_equations = [](const auto& physics)
        {
            return serif::equations::EquationRegistry{::serif::tests::equations::BarotropicClosure{model::borrow_eos(physics.equation_of_state()), 1.0}};
        }
    });
    const auto equations = model::resolve_equilibrium_equations(star);
    using Closure = ::serif::tests::equations::BarotropicClosure<model::EosReference<MoveOnlyPhysicalEOS>>;
    serif::equations::EquationPoint<Closure::Operands> point;
    point.fields.get<Closure::Test>() = 11.0;
    point.fields.get<Closure::Enthalpy>() = 2.0;
    const auto linearization = equations.interior_equations.get<Closure>().evaluate(point);
    CHECK(linearization.residual == 3.0);
    CHECK(linearization.field_derivatives.get<Closure::Enthalpy>() == -4.0);
}