#include <algorithm>
#include <concepts>
#include <memory>
#include <string_view>
#include <utility>
#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"

namespace {
    namespace model = serif::model;
    namespace backend = model::backend::mfem;
    namespace equations = serif::equations;
    namespace execution = serif::operators::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    namespace blocks = serif::discretization::blocks;
    namespace tests = serif::tests::equations;
    using EOS = serif::eos::models::Polytrope;
    using Density = blocks::density::mass;
    using Enthalpy = ::serif::tests::equations::Enthalpy;
    using Displacement = blocks::displacement::vector;
    using Kind = blocks::RepresentedValueKind;
    using Vector = mfem::Vector;

    /** @brief Require genuine ownership transfer while preserving the wrapped equation's physics and weak test. */
    template <typename Equation>
    class MoveOnlyEquation final : public equations::Contribution<equations::TestOperandOf<Equation>, equations::SupportOf<Equation>> {
    public:
        using Operands = typename Equation::Operands;
        static constexpr std::string_view name = "move-only equation";

        /** @brief Retain an evaluation counter so tests can observe lifetime without owning the equation. */
        MoveOnlyEquation(Equation equation, std::shared_ptr<int> evaluations) :
            m_equation(std::make_unique<Equation>(std::move(equation))), m_evaluations(std::move(evaluations)) {
        }

        /** @brief Forward the exact physical residual and derivatives without supplying a clone operation. */
        [[nodiscard]] auto evaluate(const equations::EquationPoint<Operands>& point) const {
            ++*m_evaluations;
            return m_equation->evaluate(point);
        }

        /** @brief Use the same integration order as the underlying physical equation. */
        template <typename Operand>
        [[nodiscard]] int additional_order(Operand operand, const int position_order) const {
            return m_equation->additional_order(operand, position_order);
        }

    private:
        std::unique_ptr<Equation> m_equation;
        std::shared_ptr<int> m_evaluations;
    };

    static_assert(!std::copy_constructible<MoveOnlyEquation<tests::Closure>>);
    static_assert(std::move_constructible<MoveOnlyEquation<tests::Closure>>);
    static_assert(serif::operators::VolumeEquation<MoveOnlyEquation<tests::Closure>>);

    /** @brief All borrowed finite-element bindings belong to the retained fixture. */
    [[nodiscard]] auto bindings(const tests::Fixture& fields) {
        return dependencies::make_dependency_context(dependencies::borrow<Density>(fields.density_binding),
            dependencies::borrow<Enthalpy>(fields.enthalpy_binding), dependencies::borrow<Displacement>(fields.displacement_binding));
    }

    /** @brief Copy evaluated residuals before perturbing the candidate state. */
    [[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 Check complete actions and exact transposes after their model and cache have gone away. */
    void check_retained(const backend::PhysicalSystem& system, const Vector& direction, const Vector& dual,
        const Vector& expected_action, const Vector& expected_transpose) {
        Vector action, transpose;
        system.jacobian().action()->Mult(direction, action);
        system.jacobian().action()->MultTranspose(dual, transpose);
        CHECK(tests::relative_error(action, expected_action) == 0.0);
        CHECK(tests::relative_error(transpose, expected_transpose) == 0.0);
    }
}

TEST_CASE("Move-only interior equations preserve evaluation derivatives and retained ownership", "[model][move_only][integration]") {
    std::weak_ptr<tests::Fixture> resources;
    std::weak_ptr<int> equation_lifetime;
    std::shared_ptr<const backend::PhysicalSystem> retained;
    Vector direction, dual, expected_action, expected_transpose, expected_residual;
    {
        auto fields = std::make_shared<tests::Fixture>();
        resources = fields;
        auto evaluations = std::make_shared<int>(0);
        equation_lifetime = evaluations;
        const EOS eos{2.0, 0.3};
        auto physical = model::make_stellar_model(eos, model::EquilibriumEquations {
            .interior_equations = equations::EquationRegistry{MoveOnlyEquation{tests::Closure{eos}, std::move(evaluations)}}
        });
        STATIC_CHECK_FALSE(std::copy_constructible<decltype(physical)>);
        auto discretization = backend::make_stellar_discretization(fields, bindings(*fields), fields->compactification, fields->mapper);
        Vector density = fields->density(), enthalpy = fields->enthalpy();
        const Vector displacement = fields->displacement(0.03);
        const auto coefficients = fields->state(density, enthalpy);
        auto evaluator = backend::bind_stellar_model(std::move(physical), std::move(discretization), coefficients, displacement);
        CHECK(*equation_lifetime.lock() == 0);
        auto evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        CHECK(*equation_lifetime.lock() > 0);
        retained = *evaluation;
        expected_residual = residual(*retained);

        tests::PreparedClosure reference(tests::Closure{eos}, fields->density_binding, fields->displacement_binding, fields->compactification, fields->mapper);
        reference.prepare(coefficients, displacement, fields->workspace);
        CHECK(tests::relative_error(expected_residual, reference.data().residual) < 2.0e-13);
        const int initial_evaluations = *equation_lifetime.lock();
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        CHECK(*evaluation == retained);
        CHECK(*equation_lifetime.lock() == initial_evaluations);

        direction = tests::deterministic(retained->jacobian().action()->Width(), 0.41);
        dual = tests::deterministic(retained->jacobian().action()->Height(), 0.83);
        retained->jacobian().action()->Mult(direction, expected_action);
        retained->jacobian().action()->MultTranspose(dual, expected_transpose);
        REQUIRE(expected_action.Norml2() > 0.0);
        CHECK((expected_action * dual) == Catch::Approx(direction * expected_transpose).margin(2.0e-12));
        const auto& layout = evaluator.layout();
        const int density_offset = static_cast<int>(layout.template offset<Kind::value, Density>());
        const int enthalpy_offset = static_cast<int>(layout.template offset<Kind::value, Enthalpy>());
        const Vector initial_density(density), initial_enthalpy(enthalpy);
        const auto perturb = [&](const double step) {
            for (int index = 0; index < density.Size(); ++index) {
                density[index] = initial_density[index] + step * direction[density_offset + index];
            }
            for (int index = 0; index < enthalpy.Size(); ++index) {
                enthalpy[index] = initial_enthalpy[index] + step * direction[enthalpy_offset + index];
            }
        };
        constexpr double step = 1.0e-6;
        perturb(step);
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        CHECK(*evaluation != retained);
        CHECK(*equation_lifetime.lock() > initial_evaluations);
        Vector difference = residual(**evaluation);
        perturb(-step);
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE(evaluation);
        difference -= residual(**evaluation);
        difference /= 2.0 * step;
        CHECK(tests::relative_error(expected_action, difference) < 2.0e-8);

        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(direction, before_rejection_action);
        before_rejection->jacobian().action()->MultTranspose(dual, before_rejection_transpose);
        enthalpy = -1.0;
        evaluation = evaluator.evaluate(coefficients, displacement);
        REQUIRE_FALSE(evaluation);
        CHECK_FALSE(evaluator.cache().ready());
        CHECK(tests::relative_error(residual(*before_rejection), before_rejection_residual) == 0.0);
        check_retained(*before_rejection, direction, dual, before_rejection_action, before_rejection_transpose);
        CHECK(tests::relative_error(residual(*retained), expected_residual) == 0.0);
        check_retained(*retained, direction, dual, expected_action, expected_transpose);
    }
    CHECK_FALSE(resources.expired());
    CHECK_FALSE(equation_lifetime.expired());
    CHECK(tests::relative_error(residual(*retained), expected_residual) == 0.0);
    check_retained(*retained, direction, dual, expected_action, expected_transpose);
    retained.reset();
    CHECK(resources.expired());
    CHECK(equation_lifetime.expired());
}