#include <algorithm>
#include <memory>
#include <string_view>
#include <vector>

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

#include "serif/backend/mfem/model/equilibrium.hpp"
#include "serif/tests/equations/fixture.hpp"
#include "serif/tests/surface/boundary_fixture.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 quantities = serif::dimensions;
    namespace tests = serif::tests::equations;
    namespace gas = serif::tests::surface;
    namespace boundary = gas::boundary;
    namespace blocks = serif::discretization::blocks;
    using Displacement = blocks::displacement::vector;
    using Kind = blocks::RepresentedValueKind;
    using Vector = mfem::Vector;

    struct Material : blocks::field {
        static constexpr std::string_view name = "test thermodynamic fields";
    };

    /** @brief Add the runtime names needed by volume quadrature to the existing boundary test terms. */
    template <typename Quantity>
    struct ScalarField : boundary::ContinuousScalar<Quantity, 3> {
        using field_type = Material;
        static constexpr std::string_view name = "test scalar thermodynamic field";
    };

    using Density = ScalarField<quantities::MassDensity>;
    using Temperature = ScalarField<quantities::Temperature>;
    using Pressure = ScalarField<quantities::Pressure>;

    /** @brief Distinguish a uniform thermodynamic state from a global constraint with coincident units. */
    struct UniformTemperature : boundary::UniformTemperature {
        using field_type = Material;
        static constexpr std::string_view name = "prescribed uniform temperature";
        static constexpr bool thermodynamic_state = true;
    };

    /** @brief Keep the interior equation independent of EOS selection so the boundary must supply its own relation. */
    template <typename Field>
    struct FieldEquation : equations::Contribution<serif::discretization::forms::Operand<Field>> {
        using Value = serif::discretization::forms::Operand<Field>;
        using Operands = serif::utils::types::TypeList<Value>;
        using ResidualQuantity = typename Field::Quantity;
        static constexpr std::string_view name = "independent scalar test equation";

        /** @brief A linear weak equation gives a nonzero interior that the surface must replace. */
        [[nodiscard]] equations::PointLinearization<Operands> evaluate(const equations::EquationPoint<Operands>& point) const {
            equations::PointLinearization<Operands> result;
            result.residual = point.fields.template get<Field>();
            result.field_derivatives.template get<Field>() = 1.0;
            return result;
        }

        /** @brief The scalar coefficient introduces no additional quadrature degree. */
        [[nodiscard]] int additional_order(Value, int) const {
            return 0;
        }
    };

    /** @brief Supply row ownership only when two different interior rows could carry a pressure residual. */
    struct DensityEquationWithSurface : FieldEquation<Density> {
        using IsobaricSurfaceField = Density;
    };

    /** @brief Expose the forward ideal-gas relation without advertising any inverse. */
    struct ForwardGasEOS : gas::GasEOS {
        using Relations = serif::eos::relations::RelationCatalog<gas::PressureFromDensityAndTemperature>;
    };

    /** @brief A pressure field already represents the boundary quantity even without numerical EOS relations. */
    struct SymbolicEOS {
    };

    struct Fixture : tests::Fixture {
        serif::discretization::backend::mfem::TermBinding<Density, tests::Schema> density{enthalpy_space};
        serif::discretization::backend::mfem::TermBinding<Temperature, tests::Schema> temperature{enthalpy_space};
        serif::discretization::backend::mfem::TermBinding<Pressure, tests::Schema> pressure{enthalpy_space};
        serif::discretization::backend::mfem::TermBinding<UniformTemperature, tests::Schema> uniform_temperature{mesh};

        /** @brief Reuse established volume resources with an independently marked x=0 stellar surface. */
        Fixture() : tests::Fixture(boundary::make_boundary_mesh()) {
        }
    };

    /** @brief Retain only the physical bindings needed by a particular model. */
    template <typename... Bindings>
    [[nodiscard]] auto discretize(const std::shared_ptr<Fixture>& fixture, Bindings... bindings) {
        return backend::make_stellar_discretization(fixture,
            dependencies::make_dependency_context(std::move(bindings)..., dependencies::borrow<Displacement>(fixture->displacement_binding)),
            fixture->compactification, fixture->mapper);
    }

    /** @brief Save a complete residual before another candidate replaces the graph's current result. */
    [[nodiscard]] Vector residual(const backend::PhysicalSystem& system) {
        Vector values(static_cast<int>(system.residual().size()));
        std::copy(system.residual().begin(), system.residual().end(), values.GetData());
        return values;
    }

    /** @brief Exercise the complete assembled action, transpose, cache reuse and retained evaluation lifetime. */
    template <typename Evaluator, typename Coefficients>
    void check_linearization(Evaluator& evaluator, const Coefficients& coefficients, const Vector& displacement, Vector& state) {
        auto evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        const auto linearization = *evaluation;
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        CHECK(*evaluation == linearization);
        CHECK(evaluator.cache().report().computed_count() == 0);

        const Vector original(state);
        const Vector direction = tests::deterministic(state.Size(), 0.6);
        const Vector dual = tests::deterministic(static_cast<int>(linearization->residual().size()), 1.1);
        Vector action(dual.Size()), transpose(state.Size());
        linearization->jacobian().action()->Mult(direction, action);
        linearization->jacobian().action()->MultTranspose(dual, transpose);
        CHECK((dual * action) == Catch::Approx(direction * transpose).margin(2.0e-12));

        constexpr double step = 1.0e-5;
        state.Add(step, direction);
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        Vector difference = residual(**evaluation);
        state = original;
        state.Add(-step, direction);
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        difference -= residual(**evaluation);
        difference /= 2.0 * step;
        for (int row = 0; row < action.Size(); ++row) {
            CHECK(action[row] == Catch::Approx(difference[row]).margin(2.0e-9));
        }
        state = original;
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        Vector frozen(action.Size());
        linearization->jacobian().action()->Mult(direction, frozen);
        for (int row = 0; row < action.Size(); ++row) {
            CHECK(frozen[row] == action[row]);
        }
    }

    /** @brief Infer a density carrier while acquiring prescribed temperature solely from the supplied physical state. */
    template <typename EOS, bool Inverse>
    void check_prescribed_temperature() {
        auto fixture = std::make_shared<Fixture>();
        const auto discretization = discretize(fixture, dependencies::borrow<Density>(fixture->density), dependencies::borrow<UniformTemperature>(fixture->uniform_temperature));
        Vector density = boundary::sequence(fixture->density.size(), 0.8, 0.3);
        Vector temperature = tests::constant(1, 1.7);
        const Vector displacement = tests::constant(fixture->displacement_binding.size(), 0.0);
        const auto coefficients = execution::make_term_coefficients(execution::bind(fixture->density, density), execution::bind(fixture->uniform_temperature, temperature));
        const auto star = model::make_stellar_model(EOS{}, model::EquilibriumEquations {
            .interior_equations = equations::EquationRegistry{FieldEquation<Density>{}},
            .shape_force = model::NoShapeForce{},
            .surface_condition = surface::Isobaric{quantities::PressureValue{0.8}}
        });
        using Definition = typename decltype(star)::Equations;
        static_assert(std::same_as<typename Definition::SurfaceCondition, surface::Isobaric>);
        auto evaluator = backend::bind_stellar_model<dependencies::KeyList<UniformTemperature>>(star, discretization, coefficients, displacement);
        static_assert(decltype(evaluator)::Plan::has_surface);
        using Formulation = typename decltype(evaluator)::SurfaceFormulation;
        STATIC_CHECK(std::same_as<typename Formulation::CarrierTerm, Density>);
        STATIC_CHECK(Formulation::resolves_carrier == Inverse);
        CHECK(evaluator.layout().template total_size<Kind::value>() == static_cast<std::size_t>(density.Size()));
        STATIC_CHECK_FALSE(dependencies::ContainsKey<UniformTemperature, typename decltype(evaluator)::Plan::Values>);
        const model::StellarState packed_state(evaluator.layout(), std::vector<double>(density.GetData(), density.GetData() + density.Size()));
        const auto prescribed = execution::make_term_coefficients(execution::bind(fixture->uniform_temperature, temperature));
        const auto initial_evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(initial_evaluation);
        using InteriorResult = execution::volume_keys::Scoped<Density>::Linearization;
        const Vector selected = boundary::boundary_mask(fixture->density);
        for (double thermal : {1.7, 2.1}) {
            const auto interior = evaluator.cache().template share<InteriorResult>();
            temperature[0] = thermal;
            const auto evaluation = evaluator.evaluate(packed_state, displacement, prescribed);
            REQUIRE(evaluation);
            CHECK(evaluator.cache().template share<InteriorResult>() == interior);
            CHECK((*evaluation)->jacobian().action()->Width() == density.Size());
            if (thermal == 1.7) {
                CHECK(*evaluation == *initial_evaluation);
                CHECK(evaluator.cache().report().computed_count() == 0);
            } else {
                CHECK(*evaluation != *initial_evaluation);
            }
            const auto external_evaluation = evaluator.evaluate(coefficients, displacement);
            REQUIRE(external_evaluation);
            CHECK(*external_evaluation == *evaluation);
            CHECK(evaluator.cache().report().computed_count() == 0);
            int boundary_rows = 0;
            for (int row = 0; row < density.Size(); ++row) {
                if (selected[row] != 0.0) {
                    ++boundary_rows;
                    const double expected = Inverse ? density[row] - 0.8 / thermal : density[row] * thermal - 0.8;
                    CHECK((*evaluation)->residual()[row] == Catch::Approx(expected).margin(2.0e-13));
                }
            }
            REQUIRE(boundary_rows > 0);
            REQUIRE(boundary_rows < density.Size());
            check_linearization(evaluator, coefficients, displacement, density);
        }
    }
} // namespace

TEST_CASE("Isobaric infers a pressure residual from a forward-only non-enthalpy EOS", "[model][surface][integration]") {
    check_prescribed_temperature<ForwardGasEOS, false>();
}

TEST_CASE("Isobaric prefers an available non-enthalpy inverse with prescribed boundary-only input", "[model][surface][integration]") {
    check_prescribed_temperature<gas::GasEOS, true>();
}

TEST_CASE("Isobaric uses an existing pressure field without requiring EOS relations", "[model][surface][integration]") {
    auto fixture = std::make_shared<Fixture>();
    const auto discretization = discretize(fixture, dependencies::borrow<Pressure>(fixture->pressure));
    Vector pressure = boundary::sequence(fixture->pressure.size(), 0.9, 0.1);
    const Vector displacement = tests::constant(fixture->displacement_binding.size(), 0.0);
    const auto coefficients = execution::make_term_coefficients(execution::bind(fixture->pressure, pressure));
    const auto star = model::make_stellar_model(SymbolicEOS{}, model::EquilibriumEquations {
        .interior_equations = equations::EquationRegistry{FieldEquation<Pressure>{}},
        .shape_force = model::NoShapeForce{},
        .surface_condition = surface::Isobaric{quantities::PressureValue{0.7}}
    });
    auto evaluator = backend::bind_stellar_model(star, discretization, coefficients, displacement);
    using Formulation = typename decltype(evaluator)::SurfaceFormulation;
    STATIC_CHECK(std::same_as<typename Formulation::CarrierTerm, Pressure>);
    const auto evaluation = evaluator.evaluate(coefficients, displacement);
    REQUIRE(evaluation);
    const Vector selected = boundary::boundary_mask(fixture->pressure);
    int boundary_rows = 0;
    for (int row = 0; row < pressure.Size(); ++row) {
        if (selected[row] != 0.0) {
            ++boundary_rows;
            CHECK((*evaluation)->residual()[row] == Catch::Approx(pressure[row] - 0.7).margin(2.0e-13));
        }
    }
    REQUIRE(boundary_rows > 0);
    check_linearization(evaluator, coefficients, displacement, pressure);
}

TEST_CASE("An equation can name its isobaric field when two continuous rows are eligible", "[model][surface][integration]") {
    auto fixture = std::make_shared<Fixture>();
    const auto discretization = discretize(fixture, dependencies::borrow<Temperature>(fixture->temperature), dependencies::borrow<Density>(fixture->density));
    const int size = fixture->density.size();
    Vector state(2 * size);
    Vector temperature(state.GetData(), size), density(state.GetData() + size, size);
    const Vector initial_temperature = boundary::sequence(size, 1.7, 0.2);
    const Vector initial_density = boundary::sequence(size, 0.8, 0.4);
    temperature = initial_temperature;
    density = initial_density;
    const Vector displacement = tests::constant(fixture->displacement_binding.size(), 0.0);
    const auto coefficients = execution::make_term_coefficients(execution::bind(fixture->temperature, temperature), execution::bind(fixture->density, density));
    const auto star = model::make_stellar_model(ForwardGasEOS{}, model::EquilibriumEquations {
        .interior_equations = equations::EquationRegistry{FieldEquation<Temperature>{}, DensityEquationWithSurface{}},
        .shape_force = model::NoShapeForce{},
        .surface_condition = surface::Isobaric{quantities::PressureValue{0.8}}
    });
    auto evaluator = backend::bind_stellar_model(star, discretization, coefficients, displacement);
    using Formulation = typename decltype(evaluator)::SurfaceFormulation;
    STATIC_CHECK(std::same_as<typename Formulation::CarrierTerm, Density>);
    CHECK(evaluator.layout().template offset<Kind::value, Temperature>() == 0);
    CHECK(evaluator.layout().template offset<Kind::value, Density>() == static_cast<std::size_t>(size));
    const auto evaluation = evaluator.evaluate(coefficients, displacement);
    REQUIRE(evaluation);
    const int row_offset = static_cast<int>(evaluator.layout().template offset<Kind::residual, Density>());
    const Vector selected = boundary::boundary_mask(fixture->density);
    for (int row = 0; row < size; ++row) {
        if (selected[row] != 0.0) {
            CHECK((*evaluation)->residual()[row_offset + row] == Catch::Approx(density[row] * temperature[row] - 0.8).margin(2.0e-13));
        }
    }
    check_linearization(evaluator, coefficients, displacement, state);
}