// MeanField normalization/spherical_harmonic_surface.cpp, composed with the existing reference geometry graph.
#include <algorithm>
#include <cmath>
#include <limits>
#include <memory>
#include <string>
#include <tuple>
#include <variant>
#include <vector>

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

#include "mfem/fem/coefficient.hpp"
#include "serif/backend/mfem/geometry/deformation/computations.hpp"
#include "serif/solver/normalization/computations.hpp"
#include "serif/tests/deformation/mesh.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace
{
    namespace dependencies = serif::utils::dependencies;
    namespace deformation = serif::geometry::deformation;
    namespace backend = deformation::backend::mfem;
    namespace normalization = serif::solver::normalization;
    namespace dimensions = serif::dimensions;
    namespace blocks = serif::discretization::blocks;
    namespace keys = deformation::keys;
    namespace normalization_keys = normalization::keys;
    using Status = dependencies::PreparationStatus;
    using Failure = std::variant<deformation::DeformationFailure, std::string>;
    using Shape = blocks::surface_deformation::coefficients;
    using Force = blocks::surface_deformation::shape_equilibrium;
    using Layout = normalization::CoordinateLayout<dependencies::KeyList<Shape>, dependencies::KeyList<Force>>;
    using Plan = normalization::NormalizationPlan<
        dependencies::Bindings<dependencies::Binding<Shape, normalization::PhysicalRieszCoordinate<dimensions::Length, keys::SurfaceMetric>>>,
        dependencies::Bindings<dependencies::Binding<Force, normalization::PhysicalRieszCoordinate<dimensions::Force, keys::SurfaceMetric>>>>;

    /** @brief Retain all borrowed reference and trial resources while selecting only reference normalization work. */
    struct SurfaceFixture
    {
        serif::tests::deformation::PairedMesh mesh{1, 0.0, 3};
        backend::ReferenceDiscretization discretization = backend::deformation_discretization(mesh.scalar_space, mesh.vector_space, *mesh.generated.reference_mesh);
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface> model{.surface = {.maximum_degree = 4}, .stellar_interior = {.radial_power = 0.5}, .vacuum = {}};
        std::vector<double> coefficients{std::numeric_limits<double>::quiet_NaN()};
        Layout layout{.state_sizes = {25}, .residual_sizes = {25}};
        normalization::StellarReferenceScales scales = *normalization::prepare_stellar_reference_scales(dimensions::QuantityValue<dimensions::Mass>{2.0}, dimensions::QuantityValue<dimensions::Length>{3.0});

        /** @brief Preserve each existing deformation resource's identity and tracking when adding normalization inputs. */
        [[nodiscard]] auto inputs() const
        {
            const auto reference = backend::make_deformation_inputs(model, discretization, coefficients);
            return dependencies::make_dependency_context(
                reference.resource<keys::ReferenceDiscretization>(),
                reference.resource<keys::SurfaceDiscretization>(),
                reference.resource<keys::ReferenceCoordinates>(),
                reference.resource<keys::SurfaceModel>(),
                reference.resource<keys::InteriorModel>(),
                reference.resource<keys::VacuumModel>(),
                reference.resource<keys::ShapeCoefficients>(),
                dependencies::borrow<normalization_keys::Layout>(layout),
                dependencies::borrow<normalization_keys::Scales>(scales));
        }
    };

    /** @brief Share one metric producer between the primal and dual coordinate factors. */
    [[nodiscard]] auto make_cache(const SurfaceFixture& fixture)
    {
        auto computations = std::tuple_cat(backend::make_deformation_computations<deformation::SphericalHarmonicSurface>(), std::tuple{backend::make_surface_metric_computation()}, 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<Failure>(std::move(graph), fixture.inputs());
    }

    /** @brief Recover MeanField's primal/dual factors and cancellation of reference Gram weights in physical work. */
    void check_harmonic_factors(const normalization::DiagonalNormalization& coordinates, const std::vector<double>& metric, const normalization::StellarReferenceScales& scales)
    {
        REQUIRE(coordinates.state_size() == metric.size());
        REQUIRE(coordinates.residual_size() == metric.size());
        double physical_work = 0.0;
        double normalized_work = 0.0;
        std::vector<double> direction(metric.size()), force(metric.size()), normalized_direction(metric.size()), normalized_force(metric.size());
        for (std::size_t mode = 0; mode < metric.size(); ++mode)
        {
            CAPTURE(mode);
            REQUIRE(std::isfinite(metric[mode]));
            REQUIRE(metric[mode] > 0.0);
            CHECK(coordinates.state_factors()[mode] == Catch::Approx(std::sqrt(metric[mode]) / scales.radius.value()).epsilon(3.0e-14));
            CHECK(coordinates.residual_factors()[mode] == Catch::Approx(1.0 / (scales.force.value() * std::sqrt(metric[mode]))).epsilon(3.0e-14));
            CHECK(coordinates.state_factors()[mode] * coordinates.residual_factors()[mode] == Catch::Approx(1.0 / (scales.radius.value() * scales.force.value())).epsilon(3.0e-14));
            direction[mode] = std::sin(0.7 * (mode + 1));
            force[mode] = std::cos(0.3 * (mode + 1));
            physical_work += direction[mode] * force[mode];
        }
        coordinates.normalize_state(direction, normalized_direction);
        coordinates.normalize_residual(force, normalized_force);
        for (std::size_t mode = 0; mode < metric.size(); ++mode)
        {
            normalized_work += normalized_direction[mode] * normalized_force[mode];
        }
        CHECK(normalized_work == Catch::Approx(physical_work / (scales.radius.value() * scales.force.value())).epsilon(5.0e-14));
    }
}

TEST_CASE("Harmonic normalization uses native shape metrics before candidate preparation", "[integration][normalization][meanfield]")
{
    SurfaceFixture fixture;
    SECTION("All real azimuthal modes")
    {
    }
    SECTION("Axisymmetric modes")
    {
        fixture.model.surface.maximum_azimuthal_order = 0;
        fixture.layout = {.state_sizes = {5}, .residual_sizes = {5}};
    }
    auto cache = make_cache(fixture);
    const auto initial = cache.try_prepare(fixture.inputs(), dependencies::RequestedResults<normalization_keys::Normalization>{});
    REQUIRE(initial);
    CHECK(initial->computed_count() == 7);
    CHECK(initial->status<keys::SurfaceMetric>() == Status::computed);
    CHECK(initial->status<keys::RadialInterpolation>() == Status::not_required);
    CHECK(initial->status<keys::VolumeExtension>() == Status::not_required);
    CHECK(initial->status<keys::SurfaceDisplacement>() == Status::not_required);
    CHECK(initial->status<keys::VolumeDisplacement>() == Status::not_required);
    CHECK_FALSE(cache.contains<keys::DomainMap>());
    CHECK_FALSE(cache.contains<keys::Linearization>());
    CHECK(cache.get<keys::SurfaceMetric>().size() == fixture.layout.state_sizes[0]);
    check_harmonic_factors(cache.get<normalization_keys::Normalization>(), cache.get<keys::SurfaceMetric>(), fixture.scales);
}

TEST_CASE("Reference changes refresh harmonic normalization while frozen solve coordinates remain unchanged", "[integration][normalization][meanfield]")
{
    SurfaceFixture fixture;
    auto cache = make_cache(fixture);
    const auto requested = dependencies::RequestedResults<normalization_keys::Normalization>{};
    REQUIRE(cache.try_prepare(fixture.inputs(), requested));
    const normalization::DiagonalNormalization frozen = cache.get<normalization_keys::Normalization>();
    const std::vector<double> original_metric = cache.get<keys::SurfaceMetric>();
    const std::vector<double> original_factors(frozen.state_factors().begin(), frozen.state_factors().end());

    fixture.coefficients.assign(25, 0.03);
    fixture.model.stellar_interior.radial_power = 3.0;
    const auto trial_changed = cache.try_prepare(fixture.inputs(), requested);
    REQUIRE(trial_changed);
    CHECK(trial_changed->computed_count() == 0);
    CHECK(trial_changed->status<keys::SurfaceMetric>() == Status::reused);
    CHECK(trial_changed->status<normalization_keys::Normalization>() == Status::reused);
    CHECK_FALSE(cache.contains<keys::RadialInterpolation>());

    fixture.scales = *normalization::prepare_stellar_reference_scales(dimensions::QuantityValue<dimensions::Mass>{4.0}, dimensions::QuantityValue<dimensions::Length>{5.0});
    const auto scales_changed = cache.try_prepare(fixture.inputs(), requested);
    REQUIRE(scales_changed);
    CHECK(scales_changed->computed_count() == 3);
    CHECK(scales_changed->status<keys::SurfaceMetric>() == Status::reused);
    CHECK(cache.get<keys::SurfaceMetric>() == original_metric);
    check_harmonic_factors(cache.get<normalization_keys::Normalization>(), original_metric, fixture.scales);
    CHECK_FALSE(std::ranges::equal(cache.get<normalization_keys::Normalization>().state_factors(), original_factors));
    CHECK(std::ranges::equal(frozen.state_factors(), original_factors));

    mfem::VectorFunctionCoefficient stretch(3, [](const mfem::Vector& position, mfem::Vector& image)
    {
        image = position;
        image[2] *= 0.85;
    });
    fixture.mesh.generated.mesh->Transform(stretch);
    const auto reference_changed = cache.try_prepare(fixture.inputs(), requested);
    REQUIRE(reference_changed);
    CHECK(reference_changed->status<keys::SurfaceNodes>() == Status::reused);
    CHECK(reference_changed->status<keys::ReferenceSurface>() == Status::computed);
    CHECK(reference_changed->status<keys::SurfaceMetric>() == Status::computed);
    CHECK(reference_changed->status<normalization_keys::Normalization>() == Status::computed);
    CHECK(cache.get<keys::SurfaceMetric>() != original_metric);
    check_harmonic_factors(cache.get<normalization_keys::Normalization>(), cache.get<keys::SurfaceMetric>(), fixture.scales);
    CHECK(std::ranges::equal(frozen.state_factors(), original_factors));
    CHECK_FALSE(cache.contains<keys::VolumeDisplacement>());
}

TEST_CASE("Prepared reference traces normalize shape without trial or extension sources", "[integration][normalization]")
{
    SurfaceFixture fixture;
    const auto reference = backend::sample_reference_surface(fixture.mesh.scalar_space, fixture.mesh.surface_attribute);
    REQUIRE(reference);
    const auto surface = deformation::build_surface_embedding(fixture.model.surface, reference->positions);
    REQUIRE(surface);
    const auto inputs = dependencies::make_dependency_context(
        dependencies::own<keys::SurfaceDiscretization>(backend::SurfaceDiscretization{.scalar_space = fixture.mesh.scalar_space, .surface_attribute = fixture.mesh.surface_attribute}, backend::SurfaceDiscretizationTracking{}),
        dependencies::own<keys::ReferenceSurface>(std::make_shared<const deformation::ReferenceSurface>(*reference)),
        dependencies::own<keys::SurfaceMap>(std::make_shared<const deformation::SurfaceEmbedding>(*surface)),
        dependencies::borrow<normalization_keys::Layout>(fixture.layout),
        dependencies::borrow<normalization_keys::Scales>(fixture.scales));
    auto computations = std::tuple_cat(std::tuple{backend::make_surface_metric_computation()}, normalization::make_normalization_computations<Plan>(fixture.layout));
    auto graph = std::apply([](auto... nodes)
    {
        return dependencies::make_computation_graph(std::move(nodes)...);
    }, std::move(computations));
    auto cache = dependencies::make_computation_cache<Failure>(std::move(graph), inputs);
    const auto prepared = cache.try_prepare(inputs);
    REQUIRE(prepared);
    CHECK(prepared->computed_count() == 4);
    check_harmonic_factors(cache.get<normalization_keys::Normalization>(), cache.get<keys::SurfaceMetric>(), fixture.scales);
}

TEST_CASE("Harmonic normalization preserves geometry and factor failure types", "[integration][normalization]")
{
    SurfaceFixture fixture;
    auto cache = make_cache(fixture);
    const auto requested = dependencies::RequestedResults<normalization_keys::Normalization>{};
    REQUIRE(cache.try_prepare(fixture.inputs(), requested));
    SECTION("Invalid harmonic geometry retains its domain failure")
    {
        fixture.model.surface.maximum_degree = -1;
        const auto failed = cache.try_prepare(fixture.inputs(), requested);
        REQUIRE_FALSE(failed);
        REQUIRE(std::holds_alternative<deformation::DeformationFailure>(failed.error().error));
        CHECK_FALSE(std::get<deformation::DeformationFailure>(failed.error().error).message.empty());
        fixture.model.surface.maximum_degree = 4;
    }
    SECTION("Incompatible normalization layout retains its factor diagnostic")
    {
        fixture.layout.state_sizes[0] = 24;
        const auto failed = cache.try_prepare(fixture.inputs(), requested);
        REQUIRE_FALSE(failed);
        REQUIRE(std::holds_alternative<std::string>(failed.error().error));
        CHECK(std::get<std::string>(failed.error().error).find("24") != std::string::npos);
        fixture.layout.state_sizes[0] = 25;
    }
    CHECK_FALSE(cache.contains<normalization_keys::Normalization>());
    REQUIRE(cache.try_prepare(fixture.inputs(), requested));
    check_harmonic_factors(cache.get<normalization_keys::Normalization>(), cache.get<keys::SurfaceMetric>(), fixture.scales);
}