#include <algorithm>
#include <cmath>
#include <limits>
#include <tuple>
#include <vector>

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

#include "mfem/fem/coefficient.hpp"
#include "serif/discretization/domain/schema/schemas.hpp"
#include "serif/backend/mfem/solver/normalization/metric_computations.hpp"
#include "serif/solver/normalization/normalization.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace
{
    namespace normalization = serif::solver::normalization;
    namespace backend = normalization::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    namespace dimensions = serif::dimensions;
    namespace discretization = serif::discretization;
    using Density = discretization::blocks::density::mass;
    using Schema = discretization::domain::schema::CoreEnvelopeVacuumDomainSchema;
    using Binding = discretization::backend::mfem::TermBinding<Density, Schema>;
    using Metric = normalization::ReferenceVolumeMetric<Density>;
    using ValueRule = normalization::PhysicalValueRule<Density>;
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<ValueRule>, dependencies::Bindings<dependencies::Binding<Density, ValueRule::Target>>>;
    using Layout = normalization::CoordinateLayout<dependencies::KeyList<Density>, dependencies::KeyList<Density>>;
    using Status = dependencies::PreparationStatus;

    /** @brief Exercise linear geometry that gains MFEM nodes during partial assembly as well as existing nodal geometry. */
    [[nodiscard]] mfem::Mesh make_reference_mesh(const bool nodal, const double length)
    {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON, 2.0 * length, length, length);
        if (nodal)
        {
            mesh.SetCurvature(2);
        }
        return mesh;
    }

    /** @brief Keep reference geometry, binding identity, and candidate coefficients independently observable. */
    struct VolumeFixture
    {
        mfem::Mesh mesh;
        mfem::H1_FECollection collection{Density::FunctionalSpace::order, 3};
        mfem::FiniteElementSpace space{&mesh, &collection};
        Binding binding{space};
        Layout layout{.state_sizes = {static_cast<std::size_t>(binding.size())}, .residual_sizes = {static_cast<std::size_t>(binding.size())}};
        normalization::StellarReferenceScales scales = *normalization::prepare_stellar_reference_scales(dimensions::QuantityValue<dimensions::Mass>{4.0}, dimensions::QuantityValue<dimensions::Length>{3.0});
        std::vector<double> candidate = std::vector<double>(binding.size(), 1.0);

        /** @brief Construct equal layouts on optionally different reference geometries. */
        explicit VolumeFixture(const bool nodal, const double length = 1.0) : mesh(make_reference_mesh(nodal, length))
        {
        }

        /** @brief Allow another live binding to replace the reference without changing context or layout types. */
        [[nodiscard]] auto inputs(Binding& reference)
        {
            return dependencies::make_dependency_context(
                backend::reference_metric_binding(reference),
                dependencies::borrow<normalization::keys::Layout>(layout),
                dependencies::borrow<normalization::keys::Scales>(scales),
                dependencies::borrow<Density>(candidate));
        }
    };

    /** @brief Compose the shared backend metric with both normalization coordinate roles. */
    [[nodiscard]] auto make_cache(VolumeFixture& fixture)
    {
        auto computations = std::tuple_cat(std::tuple{backend::make_reference_volume_metric_computation<Density, Schema>()}, normalization::make_normalization_computations<Plan>(fixture.layout));
        auto graph = std::apply([](auto... nodes)
        {
            return dependencies::make_computation_graph(std::move(nodes)...);
        }, std::move(computations));
        return dependencies::make_computation_cache<std::string>(std::move(graph), fixture.inputs(fixture.binding));
    }

    /** @brief Check exact scale cancellation independently of the graph's factor construction. */
    void check_factors(const normalization::DiagonalNormalization& factors, const std::vector<double>& metric, const double scale)
    {
        REQUIRE(factors.state_size() == metric.size());
        REQUIRE(factors.residual_size() == metric.size());
        for (std::size_t index = 0; index < metric.size(); ++index)
        {
            CHECK(factors.state_factors()[index] == Catch::Approx(std::sqrt(metric[index]) / scale).epsilon(3.0e-14));
            CHECK(factors.residual_factors()[index] == Catch::Approx(1.0 / (scale * std::sqrt(metric[index]))).epsilon(3.0e-14));
        }
    }
}

TEST_CASE("Reference volume normalization reuses one metric independently of trial coefficients", "[integration][normalization]")
{
    for (const bool nodal : {false, true})
    {
        CAPTURE(nodal);
        VolumeFixture fixture(nodal);
        CHECK((fixture.mesh.GetNodes() != nullptr) == nodal);
        auto cache = make_cache(fixture);
        const auto prepared = cache.try_prepare(fixture.inputs(fixture.binding));
        REQUIRE(prepared);
        CHECK(prepared->computed_count() == 4);
        const std::vector<double> original_metric = cache.get<Metric>();
        check_factors(cache.get<normalization::keys::Normalization>(), original_metric, fixture.scales.density.value());

        const auto unchanged = cache.try_prepare(fixture.inputs(fixture.binding));
        REQUIRE(unchanged);
        CHECK(unchanged->computed_count() == 0);
        CHECK(unchanged->status<Metric>() == Status::reused);
        std::ranges::fill(fixture.candidate, std::numeric_limits<double>::quiet_NaN());
        const auto changed_candidate = cache.try_prepare(fixture.inputs(fixture.binding));
        REQUIRE(changed_candidate);
        CHECK(changed_candidate->computed_count() == 0);
        CHECK(cache.get<Metric>() == original_metric);
    }
}

TEST_CASE("Reference volume normalization follows geometry revisions and exact binding identity", "[integration][normalization]")
{
    for (const bool nodal : {false, true})
    {
        CAPTURE(nodal);
        VolumeFixture fixture(nodal);
        auto cache = make_cache(fixture);
        REQUIRE(cache.try_prepare(fixture.inputs(fixture.binding)));
        const normalization::DiagonalNormalization frozen = cache.get<normalization::keys::Normalization>();
        const std::vector<double> original_metric = cache.get<Metric>();
        const std::vector<double> original_factors(frozen.state_factors().begin(), frozen.state_factors().end());

        mfem::VectorFunctionCoefficient dilation(3, [](const mfem::Vector& point, mfem::Vector& image)
        {
            image = point;
            image *= 2.0;
        });
        fixture.mesh.Transform(dilation);
        const auto changed_geometry = cache.try_prepare(fixture.inputs(fixture.binding));
        REQUIRE(changed_geometry);
        CHECK(changed_geometry->status<Metric>() == Status::computed);
        const std::vector<double> dilated_metric = cache.get<Metric>();
        for (std::size_t index = 0; index < original_metric.size(); ++index)
        {
            CHECK(dilated_metric[index] == Catch::Approx(8.0 * original_metric[index]).epsilon(2.0e-12));
        }
        check_factors(cache.get<normalization::keys::Normalization>(), dilated_metric, fixture.scales.density.value());
        CHECK(std::ranges::equal(frozen.state_factors(), original_factors));

        VolumeFixture replacement(nodal, 3.0);
        REQUIRE(replacement.binding.size() == fixture.binding.size());
        const auto changed_binding = cache.try_prepare(fixture.inputs(replacement.binding));
        REQUIRE(changed_binding);
        CHECK(changed_binding->status<Metric>() == Status::computed);
        const std::vector<double> replaced_metric = cache.get<Metric>();
        for (std::size_t index = 0; index < original_metric.size(); ++index)
        {
            CHECK(replaced_metric[index] == Catch::Approx(27.0 * original_metric[index]).epsilon(2.0e-12));
        }
        const auto replacement_unchanged = cache.try_prepare(fixture.inputs(replacement.binding));
        REQUIRE(replacement_unchanged);
        CHECK(replacement_unchanged->computed_count() == 0);
    }
}