#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <memory>
#include <tuple>
#include <type_traits>
#include <variant>
#include "serif/backend/mfem/model/volume_displacement.hpp"

#include "serif/tests/equations/test_equations.hpp"
#include "serif/equations/mass_normalization.hpp"
#include "serif/tests/equations/fixture.hpp"

namespace
{
    using namespace serif;
    using namespace serif::tests::equations;
    namespace equations_backend = serif::operators::backend::mfem;
    namespace dependencies = utils::dependencies;
    using Status = dependencies::PreparationStatus;
    using Mass = serif::equations::MassNormalization;

    struct FirstClosure
    {
    };
    struct SecondClosure
    {
    };
    struct MassRow
    {
    };
    using First = equations_backend::volume_keys::Scoped<FirstClosure>;
    using Second = equations_backend::volume_keys::Scoped<SecondClosure>;
    using MassStages = equations_backend::volume_keys::Scoped<MassRow>;

    /** @brief Reuse the ordinary fixture bindings without introducing an equation-specific owner. */
    auto closure_discretization(const Fixture& fixture)
    {
        return equations_backend::VolumeDiscretization{
            .residual_binding = fixture.density_binding,
            .displacement_binding = fixture.displacement_binding,
            .compactification = fixture.compactification,
            .mapper = fixture.mapper};
    }

    /** @brief Compose declarations once and retain the ordinary dependency-cache failure type. */
    template <typename Context, typename Nodes>
    auto make_cache(const Context& context, Nodes nodes)
    {
        return std::apply([&](auto... node)
                          {
                              return dependencies::make_computation_cache<equations_backend::PreparationRejection>(context, std::move(node)...);
                          },
                          std::move(nodes));
    }
} // namespace

TEST_CASE("Scoped instances of one volume equation share fields without sharing private stages", "[unit][equations][scoped_computations]")
{
    Fixture fixture;
    const auto discretization = closure_discretization(fixture);
    auto density = fixture.density();
    const auto enthalpy = fixture.enthalpy();
    const auto displacement = fixture.displacement(0.03);
    const eos::models::Polytrope first_eos{1.5, 0.25};
    const eos::models::Polytrope second_eos{3.0, 1.5};
    const auto inputs = [&]()
    {
        const auto coefficients = fixture.state(density, enthalpy);
        const auto first = equations_backend::make_volume_inputs<Closure, FirstClosure>(discretization, coefficients, displacement);
        const auto second = equations_backend::make_volume_inputs<Closure, SecondClosure>(discretization, coefficients, displacement);
        return dependencies::make_dependency_context(
            first.resource<First::Bindings>(), second.resource<Second::Bindings>(),
            first.resource<equations_backend::DisplacementCoefficients>(), first.resource<equations_backend::CompactificationCoefficients>(),
            first.resource<blocks::density::mass>(), first.resource<::serif::tests::equations::Enthalpy>());
    };
    auto cache = make_cache(inputs(), std::tuple_cat(
                                          equations_backend::make_volume_equation_computations<FirstClosure>(Closure{first_eos}, discretization),
                                          equations_backend::make_volume_equation_computations<SecondClosure>(Closure{second_eos}, discretization)));
    STATIC_CHECK_FALSE(std::same_as<First::Linearization, Second::Linearization>);
    REQUIRE(cache.try_prepare(inputs()));
    CHECK(relative_error(cache.get<First::Linearization>().residual, manual_closure(fixture, first_eos, density, enthalpy, displacement)) < 2.0e-12);
    CHECK(relative_error(cache.get<Second::Linearization>().residual, manual_closure(fixture, second_eos, density, enthalpy, displacement)) < 2.0e-12);
    const auto density_direction = deterministic(density.Size(), 0.3);
    const auto enthalpy_direction = deterministic(enthalpy.Size(), 0.7);
    const auto displacement_direction = fixture.displacement(0.1);
    const auto direction = packed(density_direction, enthalpy_direction, displacement_direction);
    const auto dual = deterministic(density.Size(), 0.9);
    mfem::Vector action, transpose;
    equations_backend::Workspace workspace;
    equations_backend::apply_jacobian(cache.get<Second::Linearization>(), direction, workspace, action);
    CHECK(relative_error(action, manual_closure(fixture, second_eos, density, enthalpy, displacement, &density_direction, &enthalpy_direction, &displacement_direction)) < 2.0e-12);
    equations_backend::apply_jacobian_transpose(cache.get<Second::Linearization>(), dual, workspace, transpose);
    CHECK((action * dual) == Catch::Approx(direction * transpose).epsilon(2.0e-12));
    const auto frozen = cache.share<Second::Linearization>();
    REQUIRE(cache.try_prepare(inputs()));
    CHECK(cache.report().computed_count() == 0);
    CHECK(cache.share<Second::Linearization>() == frozen);

    density += 0.04;
    REQUIRE(cache.try_prepare(inputs(), dependencies::RequestedResults<First::Linearization>{}));
    CHECK(cache.report().status<First::Basis<0>>() == Status::reused);
    CHECK(cache.report().status<First::QuadratureLinearization<0>>() == Status::computed);
    CHECK_FALSE(cache.contains<Second::Linearization>());
    CHECK(equations_backend::volume_work_report<FirstClosure>(cache).physics_points > 0);
    REQUIRE(cache.try_prepare(inputs()));
    CHECK(cache.report().status<First::Linearization>() == Status::reused);
    CHECK(cache.report().status<Second::QuadratureLinearization<0>>() == Status::computed);
    CHECK(cache.share<Second::Linearization>() != frozen);
    CHECK(relative_error(cache.get<Second::Linearization>().residual, manual_closure(fixture, second_eos, density, enthalpy, displacement)) < 2.0e-12);
}

TEST_CASE("Distinct volume equations prepare in one graph from common coefficient sources", "[unit][equations][scoped_computations]")
{
    Fixture fixture;
    const eos::models::Polytrope eos{1.5, 0.25};
    auto density = fixture.density();
    auto enthalpy = fixture.enthalpy();
    auto displacement = fixture.displacement(0.03);
    const auto closure = closure_discretization(fixture);
    const equations_backend::VolumeDiscretization mass{
        .residual_binding = fixture.constant_binding,
        .displacement_binding = fixture.displacement_binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper};
    const auto coefficients = [&]()
    {
        return equations_backend::make_term_coefficients(equations_backend::bind(fixture.density_binding, density), equations_backend::bind(fixture.enthalpy_binding, enthalpy));
    };
    const auto inputs = [&]()
    {
        const auto state = coefficients();
        const auto volume = equations_backend::make_volume_inputs<Closure, FirstClosure>(closure, state, displacement);
        const auto integral = equations_backend::make_volume_inputs<Mass, MassRow>(mass, state, displacement);
        return dependencies::make_dependency_context(
            volume.resource<First::Bindings>(), integral.resource<MassStages::Bindings>(),
            volume.resource<equations_backend::DisplacementCoefficients>(), volume.resource<equations_backend::CompactificationCoefficients>(),
            volume.resource<blocks::density::mass>(), volume.resource<::serif::tests::equations::Enthalpy>());
    };
    auto cache = make_cache(inputs(), std::tuple_cat(
                                          equations_backend::make_volume_equation_computations<FirstClosure>(Closure{eos}, closure),
                                          equations_backend::make_volume_equation_computations<MassRow>(Mass{0.8}, mass)));
    REQUIRE(cache.try_prepare(inputs()));
    const auto frozen_volume = cache.share<First::Linearization>();
    REQUIRE(cache.try_prepare(inputs()));
    CHECK(cache.report().computed_count() == 0);

    enthalpy += 0.03;
    REQUIRE(cache.try_prepare(inputs()));
    CHECK(cache.report().status<First::Linearization>() == Status::computed);
    CHECK(cache.report().status<MassStages::Linearization>() == Status::reused);

    density += 0.02;
    REQUIRE(cache.try_prepare(inputs()));
    CHECK(cache.report().status<MassStages::Linearization>() == Status::computed);
    enthalpy = -1.0;
    const auto rejected = cache.try_prepare(inputs());
    REQUIRE_FALSE(rejected);
    CHECK_FALSE(cache.ready());
    CHECK_FALSE(cache.contains<MassStages::Linearization>());
    CHECK(frozen_volume->residual.Size() == fixture.density_binding.size());
}

TEST_CASE("Scoped equation metadata connects the shared deformation displacement source", "[unit][equations][scoped_computations]")
{
    namespace deformation = serif::geometry::deformation;
    namespace deformation_backend = deformation::backend::mfem;
    using Failure = std::variant<equations_backend::PreparationRejection, deformation::DeformationFailure>;

    Fixture fixture;
    const auto discretization = closure_discretization(fixture);
    const auto density = fixture.density();
    const auto enthalpy = fixture.enthalpy();
    const auto displacement = fixture.displacement(0.03);
    std::vector<double> generated(displacement.GetData(), displacement.GetData() + displacement.Size());
    const auto reference = deformation_backend::deformation_discretization(fixture.enthalpy_space, fixture.displacement_space, fixture.mesh);
    const auto volume = equations_backend::make_volume_inputs<Closure, FirstClosure>(discretization, fixture.state(density, enthalpy));
    const auto context = dependencies::make_dependency_context(
        volume.resource<First::Bindings>(), volume.resource<equations_backend::CompactificationCoefficients>(),
        volume.resource<blocks::density::mass>(), volume.resource<::serif::tests::equations::Enthalpy>(),
        dependencies::own<deformation::keys::ReferenceDiscretization>(reference, deformation_backend::ReferenceDiscretizationTracking{}),
        dependencies::borrow<deformation::keys::VolumeDisplacement>(generated));
    const eos::models::Polytrope eos{1.5, 0.25};
    auto cache = std::apply([&](auto... node)
    {
        return dependencies::make_computation_cache<Failure>(context, std::move(node)...);
    }, std::tuple_cat(
        deformation_backend::make_volume_displacement_computations<First::Bindings>(fixture.displacement_binding),
        equations_backend::make_volume_equation_computations<FirstClosure>(Closure{eos}, discretization)));
    REQUIRE(cache.try_prepare(context));
    CHECK(relative_error(cache.get<equations_backend::DisplacementCoefficients>().values, displacement) == 0.0);
    CHECK(relative_error(cache.get<First::Linearization>().residual, manual_closure(fixture, eos, density, enthalpy, displacement)) < 2.0e-12);
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.report().computed_count() == 0);
    generated[0] += 0.001;
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.report().status<deformation_backend::volume_displacement_keys::CheckedBinding>() == Status::reused);
    CHECK(cache.report().status<First::Geometry<0>>() == Status::computed);
}