#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <memory>
#include <span>
#include <tuple>
#include <variant>
#include <vector>

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

#include "serif/equations/angular_momentum.hpp"
#include "serif/equations/barotropic_density.hpp"
#include "serif/equations/primal_gravity.hpp"
#include "serif/equations/surface_potential.hpp"
#include "serif/equations/mass_normalization.hpp"
#include "serif/backend/mfem/model/equilibrium.hpp"
#include "serif/backend/mfem/model/discretization.hpp"
#include "serif/backend/mfem/model/graph.hpp"
#include "serif/tests/equations/gravity_fixture.hpp"
#include "serif/tests/deformation/mesh.hpp"

namespace {
    namespace model = serif::model;
    namespace backend = model::backend::mfem;
    namespace execution = serif::operators::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    namespace equations = serif::equations;
    namespace surface = serif::geometry::surface;
    namespace deformation = serif::geometry::deformation;
    namespace blocks = serif::discretization::blocks;
    namespace test = serif::tests::equations;
    using Polytrope = serif::eos::models::Polytrope;
    using Density = blocks::density::mass;
    using Potential = blocks::gravity::potential;
    using Constant = blocks::barotropic_constant::mass_normalization;
    using Omega = blocks::rotation::angular_velocity;
    using Shape = blocks::surface_deformation::coefficients;
    using Displacement = blocks::displacement::vector;
    using Kind = blocks::RepresentedValueKind;
    using Vector = mfem::Vector;

    /** @brief Keep the physical composition explicit while exercising the public binding factory. */
    template <bool SolveOmega>
    [[nodiscard]] auto make_physical_model() {
        const Polytrope eos{2.0, 0.3};
        const auto rotation = []() {
            if constexpr (SolveOmega) {
                return std::tuple{equations::AngularMomentumConstraint{0.14}};
            } else {
                return std::tuple{};
            }
        }();
        const auto interior = std::apply([](auto... equation) {
            return equations::EquationRegistry{std::move(equation)...};
        }, std::tuple_cat(std::tuple{equations::BarotropicDensityBalance{eos, serif::dimensions::PressureValue{0.025}}, equations::PrimalPoissonBalance{2.3}, equations::MassNormalization{1.1}}, rotation));
        const auto definition = model::EquilibriumEquations {
            .interior_equations = interior,
            .shape_force = equations::SurfacePotentialBalance{},
            .surface_condition = surface::Isobaric{serif::dimensions::PressureValue{0.025}}
        };
        const auto invariants = [&]() {
            if constexpr (SolveOmega) {
                return std::tuple{model::FixedAngularMomentum{serif::dimensions::AngularMomentumValue{0.14}}};
            } else {
                return std::tuple{};
            }
        }();
        return std::apply([&](const auto&... invariant) {
            return model::make_stellar_model(eos, definition, model::FixedTotalMass{serif::dimensions::MassValue{1.1}}, invariant...);
        }, invariants);
    }

    /** @brief Keep the fixed-mass model's surface and exterior resources under one lifetime owner. */
    struct Fixture {
        stroid::StroidMesh generated{serif::tests::deformation::make_mesh()};
        test::gravity::GravityFixture fields{mfem::Mesh(*generated.mesh)};
        mfem::FiniteElementSpace reference_displacement{&fields.mesh, &fields.enthalpy_collection, 3, mfem::Ordering::byNODES};
        mfem::FiniteElementSpace compactification_space{&fields.mesh, generated.exterior_coordinate->space->FEColl()};
        mfem::GridFunction compactification{&compactification_space};

        Fixture() {
            REQUIRE(compactification.Size() == generated.exterior_coordinate->values->Size());
            compactification = *generated.exterior_coordinate->values;
        }
    };

    /** @brief Borrow bindings and the existing radial mesh extension from their retained owner. */
    [[nodiscard]] auto discretize(std::shared_ptr<Fixture> fixture) {
        const auto& fields = fixture->fields;
        const auto bindings = dependencies::make_dependency_context(
            dependencies::borrow<Density>(fields.density_binding), dependencies::borrow<Potential>(fields.potential_binding),
            dependencies::borrow<Constant>(fields.constant_binding), dependencies::borrow<Omega>(fields.angular_velocity_binding),
            dependencies::borrow<Displacement>(fields.displacement_binding));
        const auto reference = deformation::backend::mfem::deformation_discretization(
            fields.enthalpy_space, fixture->reference_displacement, *fixture->generated.reference_mesh);
        const deformation::DomainDeformation<deformation::SphericalHarmonicSurface> geometry{
            .surface = {.maximum_degree = 2}, .stellar_interior = {}, .vacuum = {}
        };
        auto discretization = backend::with_deformation(
            backend::make_stellar_discretization(fixture, bindings, fixture->compactification, fields.mapper), geometry, reference);
        return std::make_shared<const decltype(discretization)>(std::move(discretization));
    }

    /** @brief A persistent coefficient view follows the independent state ordering declared by the plan. */
    template <typename Term, typename Layout>
    [[nodiscard]] Vector field(Vector& state, const Layout& layout) {
        return Vector(state.GetData() + layout.template offset<Kind::value, Term>(), static_cast<int>(layout.template size<Kind::value, Term>()));
    }

    /** @brief Compare every residual row so scalar constraints cannot disappear behind a volume norm. */
    void check_close(const Vector& actual, const Vector& expected, const double tolerance) {
        REQUIRE(actual.Size() == expected.Size());
        for (int index = 0; index < actual.Size(); ++index) {
            CHECK(actual[index] == Catch::Approx(expected[index]).margin(tolerance));
        }
    }

    /** @brief Copy an immutable evaluated residual into independent storage for finite differences. */
    [[nodiscard]] Vector residual(const backend::PhysicalSystem& system) {
        Vector result(static_cast<int>(system.residual().size()));
        std::copy(system.residual().begin(), system.residual().end(), result.GetData());
        return result;
    }

    /** @brief A rejected candidate must expose neither stale results nor a stale work report. */
    template <typename Evaluator>
    void check_unavailable(const Evaluator& evaluator) {
        CHECK_FALSE(evaluator.cache().ready());
        CHECK_FALSE(evaluator.cache().template contains<model::keys::PhysicalSystem>());
        CHECK_THROWS_AS(evaluator.cache().template get<model::keys::PhysicalSystem>(), std::logic_error);
        CHECK_THROWS_AS(evaluator.cache().template share<model::keys::PhysicalSystem>(), std::logic_error);
        CHECK(evaluator.cache().report().computed_count() == 0);
        CHECK(evaluator.cache().report().reused_count() == 0);
    }

    /** @brief Check physical assembly and candidate publication through the public stellar-model interface. */
    template <bool SolveOmega>
    void check_equilibrium_graph() {
        std::weak_ptr<Fixture> resources;
        std::shared_ptr<const backend::PhysicalSystem> retained;
        Vector retained_direction, retained_action;
        {
            auto fixture = std::make_shared<Fixture>();
            resources = fixture;
            const auto discretization = discretize(fixture);
            const auto physical = std::make_shared<const decltype(make_physical_model<SolveOmega>())>(make_physical_model<SolveOmega>());
            using Description = std::remove_cvref_t<decltype(physical->equations())>;
            using Prescribed = std::conditional_t<SolveOmega, dependencies::KeyList<>, dependencies::KeyList<Omega>>;
            using Plan = model::EquilibriumPlan<Polytrope, Description, Prescribed>;
            const std::array initial_shape{0.004, -0.002, 0.001, 0.003, -0.001, 0.002, -0.003, 0.001, 0.002};
            const auto layout = backend::equilibrium_layout<Plan>(*discretization, initial_shape);
            REQUIRE(layout.template total_size<Kind::value>() == layout.template total_size<Kind::residual>());
            STATIC_CHECK_FALSE(dependencies::ContainsKey<::serif::tests::equations::Enthalpy, typename Plan::Values>);
            const int size = static_cast<int>(layout.template total_size<Kind::value>());
            Vector state(size);
            Vector density = field<Density>(state, layout);
            Vector potential = field<Potential>(state, layout);
            Vector constant = field<Constant>(state, layout);
            Vector omega = [&]() {
                if constexpr (SolveOmega) {
                    return field<Omega>(state, layout);
                } else {
                    return test::constant(1, 0.31);
                }
            }();
            Vector shape = field<Shape>(state, layout);
            REQUIRE(static_cast<std::size_t>(shape.Size()) == initial_shape.size());
            std::copy(initial_shape.begin(), initial_shape.end(), shape.GetData());
            const auto geometry = [&]() { return std::span<const double>{shape.GetData(), static_cast<std::size_t>(shape.Size())}; };
            // Copy from lvalues: moving an MFEM vector would detach these views from the packed state.
            const Vector initial_density = fixture->fields.density();
            const Vector initial_potential = fixture->fields.potential();
            density = initial_density;
            potential = initial_potential;
            constant = 0.72;
            omega = 0.31;
            const auto make_coefficients = [&](const auto& density_binding, const auto& omega_binding) {
                return execution::make_term_coefficients(
                    execution::bind(density_binding, density), execution::bind(fixture->fields.potential_binding, potential),
                    execution::bind(fixture->fields.constant_binding, constant), execution::bind(omega_binding, omega));
            };
            const auto coefficients = make_coefficients(fixture->fields.density_binding, fixture->fields.angular_velocity_binding);
            const auto bound = backend::bind_model_definition<Prescribed>(physical, discretization, coefficients);
            REQUIRE(bound->boundary->selected_rows->size() > 0);
            REQUIRE(bound->boundary->selected_rows->size() < density.Size());
            auto evaluator = backend::bind_stellar_model<Prescribed>(*physical, *discretization, coefficients, geometry());
            check_unavailable(evaluator);

            const double initial_value = density[0];
            density[0] = std::numeric_limits<double>::quiet_NaN();
            REQUIRE_FALSE(evaluator.evaluate(coefficients, geometry()));
            check_unavailable(evaluator);
            CHECK(evaluator.cache().preparation_count() == 0);
            density[0] = initial_value;

            auto evaluation = evaluator.evaluate(coefficients, geometry());
            REQUIRE(evaluation);
            const auto initial_system = *evaluation;
            const Vector original_state(state);
            const Vector original_residual = residual(*initial_system);
            evaluation = evaluator.evaluate(coefficients, geometry());
            REQUIRE(evaluation);
            CHECK(*evaluation == initial_system);
            CHECK(evaluator.cache().report().computed_count() == 0);

            const Vector frozen_direction = test::deterministic(size, 0.8);
            Vector frozen_action, frozen_transpose;
            initial_system->jacobian().action()->Mult(frozen_direction, frozen_action);
            initial_system->jacobian().action()->MultTranspose(frozen_direction, frozen_transpose);

            // Changing input representation must not create a new physical candidate or invalidate reusable work.
            const model::StellarState packed_state(layout, std::vector<double>(state.GetData(), state.GetData() + state.Size()));
            const auto prescribed = [&]() {
                if constexpr (SolveOmega) {
                    return dependencies::make_dependency_context();
                } else {
                    return execution::make_term_coefficients(execution::bind(fixture->fields.angular_velocity_binding, omega));
                }
            }();
            STATIC_CHECK(dependencies::ContainsKey<Omega, typename decltype(evaluator)::Plan::Values> == SolveOmega);
            evaluation = evaluator.evaluate(packed_state, prescribed);
            REQUIRE(evaluation);
            CHECK(*evaluation == initial_system);
            CHECK(evaluator.cache().report().computed_count() == 0);
            evaluation = evaluator.evaluate(coefficients, geometry());
            REQUIRE(evaluation);
            CHECK(*evaluation == initial_system);
            CHECK(evaluator.cache().report().computed_count() == 0);

            // Identical compile-time coordinates can still describe incompatible meshes or equation sizes.
            for (const bool wrong_state_size : {true, false}) {
                CAPTURE(SolveOmega, wrong_state_size);
                auto incompatible_layout = layout;
                if (wrong_state_size) {
                    ++incompatible_layout.state_sizes[0];
                } else {
                    ++incompatible_layout.residual_sizes[0];
                }
                const model::StellarState incompatible_state(incompatible_layout);
                const auto preparations = evaluator.cache().preparation_count();
                CHECK_THROWS_WITH(evaluator.evaluate(incompatible_state, prescribed), Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT"));
                check_unavailable(evaluator);
                CHECK(evaluator.cache().preparation_count() == preparations);
                check_close(residual(*initial_system), original_residual, 0.0);
                Vector action, transpose;
                initial_system->jacobian().action()->Mult(frozen_direction, action);
                initial_system->jacobian().action()->MultTranspose(frozen_direction, transpose);
                check_close(action, frozen_action, 0.0);
                check_close(transpose, frozen_transpose, 0.0);
                evaluation = evaluator.evaluate(packed_state, prescribed);
                REQUIRE(evaluation);
                CHECK(*evaluation == initial_system);
                CHECK(evaluator.cache().report().computed_count() == 0);
            }

            const std::array fields{&density, &potential, &constant, &omega};
            for (std::size_t field_index = 0; field_index < fields.size(); ++field_index) {
                Vector& values = *fields[field_index];
                const double original_value = values[0];
                for (double invalid : {std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::infinity(), -std::numeric_limits<double>::infinity()}) {
                    CAPTURE(SolveOmega, field_index, invalid);
                    const auto preparations = evaluator.cache().preparation_count();
                    values[0] = invalid;
                    const auto rejected = evaluator.evaluate(coefficients, geometry());
                    REQUIRE_FALSE(rejected);
                    REQUIRE(std::holds_alternative<execution::PreparationRejection>(rejected.error().error));
                    CHECK(std::get<execution::PreparationRejection>(rejected.error().error).reason == execution::PreparationRejectionReason::NON_FINITE);
                    check_unavailable(evaluator);
                    CHECK(evaluator.cache().preparation_count() == preparations);
                    check_close(residual(*initial_system), original_residual, 0.0);
                    Vector action, transpose;
                    initial_system->jacobian().action()->Mult(frozen_direction, action);
                    initial_system->jacobian().action()->MultTranspose(frozen_direction, transpose);
                    check_close(action, frozen_action, 0.0);
                    check_close(transpose, frozen_transpose, 0.0);

                    values[0] = original_value;
                    evaluation = evaluator.evaluate(coefficients, geometry());
                    REQUIRE(evaluation);
                    CHECK(*evaluation == initial_system);
                    CHECK(evaluator.cache().report().computed_count() == 0);
                }
            }

            // Geometry failures originate in mapping, after the input context has been built.
            for (Vector* values : {&shape, static_cast<Vector*>(&fixture->compactification)}) {
                INFO((values == &shape ? "Invalid surface shape" : "Invalid compactification"));
                const Vector original_values(*values);
                *values = std::numeric_limits<double>::quiet_NaN();
                const auto preparations = evaluator.cache().preparation_count();
                REQUIRE_FALSE(evaluator.evaluate(coefficients, geometry()));
                check_unavailable(evaluator);
                CHECK(evaluator.cache().preparation_count() == preparations);
                check_close(residual(*initial_system), original_residual, 0.0);
                *values = original_values;
                evaluation = evaluator.evaluate(coefficients, geometry());
                REQUIRE(evaluation);
                CHECK(*evaluation == initial_system);
                CHECK(evaluator.cache().report().computed_count() == 0);
            }

            // Invalid sizes remain caller errors, but they must still clear preparation availability.
            Vector wrong_size(density.Size() + 1);
            wrong_size = 1.0;
            density.Swap(wrong_size);
            CHECK_THROWS_AS(evaluator.evaluate(coefficients, geometry()), std::invalid_argument);
            check_unavailable(evaluator);
            check_close(residual(*initial_system), original_residual, 0.0);
            density.Swap(wrong_size);
            evaluation = evaluator.evaluate(coefficients, geometry());
            REQUIRE(evaluation);
            CHECK(evaluator.cache().report().computed_count() == 0);

            // Identical meshes and coefficient counts do not give an external binding the owner's lifetime.
            const serif::discretization::backend::mfem::TermBinding<Density, test::Schema> other_density(fixture->fields.density_space);
            const serif::discretization::backend::mfem::TermBinding<Omega, test::Schema> other_omega(fixture->fields.mesh);
            const auto foreign_density = make_coefficients(other_density, fixture->fields.angular_velocity_binding);
            const auto foreign_omega = make_coefficients(fixture->fields.density_binding, other_omega);
            for (const auto& foreign : {foreign_density, foreign_omega}) {
                CHECK_THROWS_WITH(backend::bind_stellar_model<Prescribed>(*physical, *discretization, foreign, geometry()), Catch::Matchers::ContainsSubstring("CACHE_FE_BINDING_ERROR"));
                CHECK_THROWS_WITH(evaluator.evaluate(foreign, geometry()), Catch::Matchers::ContainsSubstring("CACHE_FE_BINDING_ERROR"));
                check_unavailable(evaluator);
                check_close(residual(*initial_system), original_residual, 0.0);
                evaluation = evaluator.evaluate(coefficients, geometry());
                REQUIRE(evaluation);
                CHECK(evaluator.cache().report().computed_count() == 0);
            }

            const int density_row = static_cast<int>(layout.template offset<Kind::residual, Density>());
            const double boundary_density = serif::eos::evaluate<serif::dimensions::MassDensity>(
                physical->equation_of_state(), serif::dimensions::PressureValue{0.025}).value();
            for (const int row : bound->boundary->selected_rows->reduced_dofs()) {
                CHECK(initial_system->residual()[density_row + row] == Catch::Approx(density[row] - boundary_density).margin(2.0e-14));
            }
            const auto infinity = serif::discretization::backend::mfem::select_boundary_dofs<
                serif::discretization::domain::InfinitySurfaceBoundary>(fixture->fields.potential_binding);
            REQUIRE(infinity.size() > 0);
            const int potential_row = static_cast<int>(layout.template offset<Kind::residual, Potential>());
            for (const int row : infinity.reduced_dofs()) {
                CHECK(initial_system->residual()[potential_row + row] == Catch::Approx(potential[row]).margin(2.0e-14));
            }

            const Vector dual = test::deterministic(size, 1.7);
            Vector transpose(size);
            initial_system->jacobian().action()->MultTranspose(dual, transpose);
            for (int variation = 0; variation < (SolveOmega ? 4 : 3); ++variation) {
                CAPTURE(variation);
                Vector direction = test::deterministic(size, 0.4);
                if (variation > 0) {
                    direction = 0.0;
                    std::size_t column = variation == 1 ? layout.template offset<Kind::value, Constant>() : layout.template offset<Kind::value, Shape>();
                    if constexpr (SolveOmega) {
                        if (variation == 3) {
                            column = layout.template offset<Kind::value, Omega>();
                        }
                    }
                    direction[static_cast<int>(column)] = 0.13;
                }
                Vector action(size);
                initial_system->jacobian().action()->Mult(direction, action);
                CHECK((action * dual) == Catch::Approx(direction * transpose).margin(2.0e-12));
                // Tiny geometry differences lose digits in the exterior residual. Cancel the leading
                // central-difference truncation term while keeping the original absolute error bound.
                const auto centered_difference = [&](const double step) {
                    state = original_state;
                    state.Add(step, direction);
                    evaluation = evaluator.evaluate(coefficients, geometry());
                    REQUIRE(evaluation);
                    Vector difference = residual(**evaluation);
                    state = original_state;
                    state.Add(-step, direction);
                    evaluation = evaluator.evaluate(coefficients, geometry());
                    REQUIRE(evaluation);
                    difference -= residual(**evaluation);
                    difference /= 2.0 * step;
                    return difference;
                };
                const double step = 1.0e-3;
                const Vector coarse = centered_difference(step);
                Vector difference = centered_difference(0.5 * step);
                difference *= 4.0;
                difference -= coarse;
                difference /= 3.0;
                check_close(action, difference, 2.0e-9);
                state = original_state;
                evaluation = evaluator.evaluate(coefficients, geometry());
                REQUIRE(evaluation);
                check_close(residual(**evaluation), original_residual, 2.0e-14);
                Vector frozen_action(size);
                initial_system->jacobian().action()->Mult(direction, frozen_action);
                check_close(frozen_action, action, 0.0);
            }

            using PoissonResult = execution::volume_keys::Scoped<Potential>::Linearization;
            using MomentumResult = execution::volume_keys::Scoped<Density>::Linearization;
            const auto poisson = evaluator.cache().template share<PoissonResult>();
            const auto momentum = evaluator.cache().template share<MomentumResult>();
            omega[0] += 0.02;
            if constexpr (SolveOmega) {
                evaluation = evaluator.evaluate(coefficients, geometry());
            } else {
                evaluation = evaluator.evaluate(packed_state, prescribed);
            }
            REQUIRE(evaluation);
            CHECK((*evaluation)->jacobian().action()->Width() == size);
            if constexpr (!SolveOmega) {
                const auto external_evaluation = evaluator.evaluate(coefficients, geometry());
                REQUIRE(external_evaluation);
                CHECK(*external_evaluation == *evaluation);
                CHECK(evaluator.cache().report().computed_count() == 0);
            }
            CHECK(evaluator.cache().template share<PoissonResult>() == poisson);
            CHECK(evaluator.cache().template share<MomentumResult>() != momentum);
            CHECK(*evaluation != initial_system);
            check_close(residual(*initial_system), original_residual, 0.0);

            const auto before_rejection = *evaluation;
            const Vector before_rejection_residual = residual(*before_rejection);
            Vector before_rejection_action, before_rejection_transpose;
            before_rejection->jacobian().action()->Mult(frozen_direction, before_rejection_action);
            before_rejection->jacobian().action()->MultTranspose(frozen_direction, before_rejection_transpose);
            density = -1.0;
            const auto rejected = evaluator.evaluate(coefficients, geometry());
            REQUIRE_FALSE(rejected);
            REQUIRE(std::holds_alternative<execution::PreparationRejection>(rejected.error().error));
            CHECK(std::get<execution::PreparationRejection>(rejected.error().error).reason == execution::PreparationRejectionReason::EQUATION);
            check_unavailable(evaluator);
            check_close(residual(*before_rejection), before_rejection_residual, 0.0);
            Vector rejected_action, rejected_transpose;
            before_rejection->jacobian().action()->Mult(frozen_direction, rejected_action);
            before_rejection->jacobian().action()->MultTranspose(frozen_direction, rejected_transpose);
            check_close(rejected_action, before_rejection_action, 0.0);
            check_close(rejected_transpose, before_rejection_transpose, 0.0);
            check_close(residual(*initial_system), original_residual, 0.0);
            state = original_state;
            omega = 0.31;
            evaluation = evaluator.evaluate(coefficients, geometry());
            REQUIRE(evaluation);
            check_close(residual(**evaluation), original_residual, 2.0e-14);
            retained = *evaluation;
            retained_direction = test::deterministic(size, 0.8);
            retained->jacobian().action()->Mult(retained_direction, retained_action);
        }
        CHECK_FALSE(resources.expired());
        Vector after_owner_release;
        retained->jacobian().action()->Mult(retained_direction, after_owner_release);
        check_close(after_owner_release, retained_action, 0.0);
        retained.reset();
        CHECK(resources.expired());
    }
} // namespace

TEST_CASE("Equilibrium graph combines the density balance, primal gravity and moving isobar rows", "[unit][model][equilibrium_graph]") {
    check_equilibrium_graph<true>();
}

TEST_CASE("Prescribed rotation remains a physical input without a Newton column or angular momentum equation", "[unit][model][equilibrium_graph]") {
    check_equilibrium_graph<false>();
}