#include <array>
#include <cmath>
#include <limits>
#include <span>
#include <tuple>
#include <utility>
#include <vector>

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

#include "mfem/fem/gridfunc.hpp"
#include "mfem/fem/linearform.hpp"
#include "serif/tests/equations/test_equations.hpp"
#include "serif/backend/mfem/geometry/deformation/reference_geometry.hpp"
#include "serif/backend/mfem/geometry/deformation/surface_metric.hpp"
#include "serif/backend/mfem/solver/normalization/metric_computations.hpp"
#include "serif/solver/normalization/normalization.hpp"
#include "serif/tests/deformation/mesh.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace {
    namespace normalization = serif::solver::normalization;
    namespace dimensions = serif::dimensions;
    namespace dependencies = serif::utils::dependencies;
    namespace discretization = serif::discretization;
    namespace blocks = discretization::blocks;
    namespace deformation = serif::geometry::deformation;
    using Schema = discretization::domain::schema::CoreEnvelopeVacuumDomainSchema;
    using Kind = blocks::RepresentedValueKind;

    /** @brief Represent the same reference star after changing numerical mass, length, and time units. */
    [[nodiscard]] normalization::StellarReferenceScales reference_scales(double length, double mass, double time) {
        const auto result = normalization::prepare_stellar_reference_scales(
            dimensions::QuantityValue<dimensions::Mass>{2.0 * mass},
            dimensions::QuantityValue<dimensions::Length>{3.0 * length},
            dimensions::QuantityValue<dimensions::GravitationalConstant>{0.7 * length * length * length / (mass * time * time)});
        REQUIRE(result);
        return *result;
    }

    /** @brief Exercise the public dependency computations with independent physical state and residual quantities. */
    template <typename Value, typename Residual, typename ValueQuantity, typename ResidualQuantity, typename Metric>
    [[nodiscard]] normalization::DiagonalNormalization prepare_factors(const std::vector<double>& metric, const normalization::StellarReferenceScales& scales) {
        using Layout = normalization::CoordinateLayout<dependencies::KeyList<Value>, dependencies::KeyList<Residual>>;
        using Plan = normalization::NormalizationPlan<
            dependencies::Bindings<dependencies::Binding<Value, normalization::DimensionlessRieszCoordinate<ValueQuantity, Metric>>>,
            dependencies::Bindings<dependencies::Binding<Residual, normalization::DimensionlessRieszCoordinate<ResidualQuantity, Metric>>>>;
        const Layout layout{.state_sizes = {metric.size()}, .residual_sizes = {metric.size()}};
        const auto inputs = dependencies::make_dependency_context(
            dependencies::borrow<normalization::keys::Layout>(layout),
            dependencies::borrow<normalization::keys::Scales>(scales), dependencies::borrow<Metric>(metric));
        auto graph = std::apply([](auto... nodes) {
            return dependencies::make_computation_graph(std::move(nodes)...);
        }, normalization::make_normalization_computations<Plan>(layout));
        auto cache = dependencies::make_computation_cache<std::string>(std::move(graph), inputs);
        REQUIRE(cache.try_prepare(inputs));
        return cache.template get<normalization::keys::Normalization>();
    }

    /** @brief Keep dimensionless vectors and their measured norms together for cross-unit comparisons. */
    struct MeasuredCoordinates {
        std::vector<double> state;
        std::vector<double> residual;
        double state_norm_squared;
        double residual_norm_squared;
    };

    /** @brief Compare actual coordinates, because equal norms alone could hide different block weights or permutations. */
    void check_same_coordinates(const MeasuredCoordinates& first, const MeasuredCoordinates& second) {
        REQUIRE(first.state.size() == second.state.size());
        REQUIRE(first.residual.size() == second.residual.size());
        for (std::size_t index = 0; index < first.state.size(); ++index) {
            CHECK(second.state[index] == Catch::Approx(first.state[index]).epsilon(2.0e-11).margin(1.0e-13));
            CHECK(second.residual[index] == Catch::Approx(first.residual[index]).epsilon(2.0e-11).margin(1.0e-13));
        }
        CHECK(second.state_norm_squared == Catch::Approx(first.state_norm_squared).epsilon(2.0e-11));
        CHECK(second.residual_norm_squared == Catch::Approx(first.residual_norm_squared).epsilon(2.0e-11));
    }

    /** @brief Use real MFEM projection and weak integration of scalar fields. */
    template <typename Term>
    [[nodiscard]] MeasuredCoordinates volume_coordinates(mfem::FiniteElementCollection& collection, double length, double mass, double time) {
        mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON, 2.0 * length, 1.5 * length, 0.75 * length);
        for (int element = 0; element < mesh.GetNE(); ++element) {
            mesh.GetElement(element)->SetAttribute(2);
        }
        mesh.SetAttributes();
        mfem::FiniteElementSpace space(&mesh, &collection);
        const discretization::backend::mfem::TermBinding<Term, Schema> binding{space};
        const auto metric = discretization::backend::mfem::reference_mass_diagonal(binding);
        REQUIRE(metric);
        const normalization::StellarReferenceScales scales = reference_scales(length, mass, time);
        const double physical_scale = normalization::characteristic_scale(scales, typename Term::Quantity{}).value();
        mfem::GridFunction field(&space);
        mfem::LinearForm residual(&space);
        mfem::ConstantCoefficient scalar(0.4 * physical_scale);
        field.ProjectCoefficient(scalar);
        residual.AddDomainIntegrator(new mfem::DomainLFIntegrator(scalar));
        residual.Assemble();
        const mfem::Vector state_coefficients = binding.gather(field);
        const mfem::Vector residual_coefficients = binding.gather(residual);
        const std::span<const double> state{state_coefficients.GetData(), static_cast<std::size_t>(state_coefficients.Size())};
        const std::span<const double> weak{residual_coefficients.GetData(), static_cast<std::size_t>(residual_coefficients.Size())};
        const auto factors = prepare_factors<Term, Term, typename Term::Quantity, typename Term::Quantity, normalization::ReferenceVolumeMetric<Term>>(*metric, scales);
        MeasuredCoordinates result{.state = std::vector<double>(state.size()), .residual = std::vector<double>(weak.size()),
            .state_norm_squared = factors.state_norm_squared(state), .residual_norm_squared = factors.residual_norm_squared(weak)};
        factors.normalize_state(state, result.state);
        factors.normalize_residual(weak, result.residual);
        return result;
    }

    /** @brief Supply an exact cellwise norm independent of a diagonal approximation to a higher-order Gram matrix. */
    struct ConstantDensityField final : blocks::field {
        struct Values final : blocks::term<ConstantDensityField> {
            using FunctionalSpace = discretization::functional_space::DiscreteFunctionalSpace<discretization::functional_space::L2, 0, 0>;
            using Support = discretization::domain::StellarDomains;
            using Quantity = dimensions::MassDensity;
        };
    };

    /** @brief Build each native surface metric on the same physical sphere expressed in different length units. */
    template <typename Surface>
    [[nodiscard]] MeasuredCoordinates surface_coordinates(const Surface& surface, double length, double mass, double time) {
        serif::tests::deformation::PairedMesh mesh(0, 0.0, 2);
        mfem::VectorFunctionCoefficient dilation(3, [=](const mfem::Vector& point, mfem::Vector& image) {
            image = point;
            image *= length;
        });
        mesh.generated.mesh->Transform(dilation);
        const auto reference = deformation::backend::mfem::sample_reference_surface(mesh.scalar_space, mesh.surface_attribute);
        REQUIRE(reference);
        const auto embedding = deformation::build_surface_embedding(surface, reference->positions);
        REQUIRE(embedding);
        const auto metric = deformation::backend::mfem::surface_parameter_mass_diagonal(mesh.scalar_space, mesh.surface_attribute, *reference, *embedding);
        REQUIRE(metric);
        const normalization::StellarReferenceScales scales = reference_scales(length, mass, time);
        const auto factors = prepare_factors<blocks::surface_deformation::coefficients, blocks::surface_deformation::shape_equilibrium,
            dimensions::Length, dimensions::Force, deformation::keys::SurfaceMetric>(*metric, scales);
        std::vector<double> state(metric->size()), residual(metric->size());
        double expected_state_norm = 0.0;
        double expected_residual_norm = 0.0;
        double expected_work = 0.0;
        for (std::size_t index = 0; index < metric->size(); ++index) {
            const double shape = std::sin(0.7 * (index + 1));
            const double force = std::cos(0.3 * (index + 1));
            state[index] = shape * scales.radius.value();
            residual[index] = force * scales.force.value();
            const double area_fraction = (*metric)[index] / std::pow(scales.radius.value(), 2);
            expected_state_norm += area_fraction * shape * shape;
            expected_residual_norm += force * force / area_fraction;
            expected_work += shape * force;
        }
        MeasuredCoordinates result{.state = std::vector<double>(metric->size()), .residual = std::vector<double>(metric->size()),
            .state_norm_squared = factors.state_norm_squared(state), .residual_norm_squared = factors.residual_norm_squared(residual)};
        factors.normalize_state(state, result.state);
        factors.normalize_residual(residual, result.residual);
        CHECK(result.state_norm_squared == Catch::Approx(expected_state_norm).epsilon(2.0e-13));
        CHECK(result.residual_norm_squared == Catch::Approx(expected_residual_norm).epsilon(2.0e-13));
        double work = 0.0;
        for (std::size_t index = 0; index < metric->size(); ++index) {
            work += result.state[index] * result.residual[index];
        }
        CHECK(work == Catch::Approx(expected_work).epsilon(2.0e-12));
        return result;
    }
}

TEST_CASE("Dimensionless volume coordinates are independent of mass length and time units", "[normalization][dimensionless][integration]") {
    mfem::H1_FECollection enthalpy(3, 3);
    mfem::H1_FECollection density(blocks::density::mass::FunctionalSpace::order, 3);
    const auto check = []<typename Term>(mfem::FiniteElementCollection& collection) {
        const MeasuredCoordinates baseline = volume_coordinates<Term>(collection, 1.0, 1.0, 1.0);
        for (const std::array<double, 3> conversion : {std::array{100.0, 1000.0, 1.0}, std::array{0.03, 7.0, 0.2}}) {
            check_same_coordinates(baseline, volume_coordinates<Term>(collection, conversion[0], conversion[1], conversion[2]));
        }
    };
    check.template operator()<::serif::tests::equations::Enthalpy>(enthalpy);
    check.template operator()<blocks::density::mass>(density);
}

TEST_CASE("Dimensionless cellwise state and weak residual norms recover the volume integral", "[normalization][dimensionless][math]") {
    mfem::L2_FECollection collection(0, 3);
    for (const double length : {1.0, 100.0}) {
        const MeasuredCoordinates result = volume_coordinates<ConstantDensityField::Values>(collection, length, 1000.0, 0.2);
        const double integral = 0.4 * 0.4 * (2.0 * 1.5 * 0.75) / std::pow(3.0, 3);
        CHECK(result.state_norm_squared == Catch::Approx(integral).epsilon(2.0e-13));
        CHECK(result.residual_norm_squared == Catch::Approx(integral).epsilon(2.0e-13));
    }
}

TEST_CASE("Dimensionless harmonic and nodal shape norms retain force work across unit changes", "[normalization][dimensionless][integration]") {
    const auto check = []<typename Surface>(const Surface& surface) {
        const MeasuredCoordinates baseline = surface_coordinates(surface, 1.0, 1.0, 1.0);
        check_same_coordinates(baseline, surface_coordinates(surface, 100.0, 1000.0, 0.2));
    };
    check(deformation::SphericalHarmonicSurface{.maximum_degree = 2});
    check(deformation::NodalRadialSurface{});
}

TEST_CASE("Dimensionless point replacements retain their own pressure scale", "[normalization][dimensionless][math]") {
    using Enthalpy = ::serif::tests::equations::Enthalpy;
    using Metric = normalization::ReferenceVolumeMetric<Enthalpy>;
    struct Selection {};
    using Layout = normalization::CoordinateLayout<dependencies::KeyList<>, dependencies::KeyList<Enthalpy>>;
    using Method = normalization::WithPointRows<normalization::DimensionlessRieszCoordinate<dimensions::SpecificEnthalpy, Metric>, dimensions::Pressure, Selection>;
    using Plan = normalization::NormalizationPlan<dependencies::Bindings<>, dependencies::Bindings<dependencies::Binding<Enthalpy, Method>>>;
    const Layout layout{.state_sizes = {}, .residual_sizes = {2}};
    const std::vector<int> rows{1};
    for (const double length : {1.0, 100.0}) {
        const auto scales = reference_scales(length, 1000.0, 0.2);
        const std::vector<double> metric{4.0 * length * length * length, std::numeric_limits<double>::quiet_NaN()};
        const auto inputs = dependencies::make_dependency_context(dependencies::borrow<normalization::keys::Layout>(layout),
            dependencies::borrow<normalization::keys::Scales>(scales), dependencies::borrow<Metric>(metric), dependencies::borrow<Selection>(rows));
        auto graph = std::apply([](auto... nodes) {
            return dependencies::make_computation_graph(std::move(nodes)...);
        }, normalization::make_normalization_computations<Plan>(layout));
        auto cache = dependencies::make_computation_cache<std::string>(std::move(graph), inputs);
        REQUIRE(cache.try_prepare(inputs));
        const auto& factors = cache.get<normalization::keys::Normalization>();
        std::array<double, 2> residual{scales.specific_energy.value() * 8.0 * std::pow(length, 3), scales.pressure.value() * 0.7};
        std::array<double, 2> normalized;
        factors.normalize_residual(residual, normalized);
        CHECK(normalized[0] == Catch::Approx(4.0 / std::sqrt(27.0)).epsilon(2.0e-13));
        CHECK(normalized[1] == Catch::Approx(0.7).epsilon(2.0e-13));
    }
}

TEST_CASE("Dimensionless geometry preparation reports invalid or unrepresentable scales", "[normalization][dimensionless][validation]") {
    for (const double invalid : {0.0, -1.0, std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()}) {
        CHECK_FALSE(normalization::detail::dimensionless_riesz_scale(Kind::value, normalization::ReferenceGeometry::volume, 1.0, invalid));
        CHECK_FALSE(normalization::detail::dimensionless_riesz_scale(Kind::residual, normalization::ReferenceGeometry::volume, invalid, 1.0));
    }
    CHECK_FALSE(normalization::detail::dimensionless_riesz_scale(Kind::value, normalization::ReferenceGeometry::volume, 1.0e300, 1.0e100));
    CHECK_FALSE(normalization::detail::dimensionless_riesz_scale(Kind::value, static_cast<normalization::ReferenceGeometry>(99), 1.0, 1.0));
    CHECK_FALSE(normalization::detail::dimensionless_riesz_scale(static_cast<Kind>(99), normalization::ReferenceGeometry::volume, 1.0, 1.0));
}

TEST_CASE("Dimensionless global coordinates retain independent unknown and constraint scales", "[normalization][dimensionless][math]") {
    using MassConstant = blocks::barotropic_constant::mass_normalization;
    for (const std::array<double, 3> conversion : {std::array{1.0, 1.0, 1.0}, std::array{100.0, 1000.0, 0.2}}) {
        const auto scales = reference_scales(conversion[0], conversion[1], conversion[2]);
        const auto factors = prepare_factors<MassConstant, MassConstant, dimensions::GravitationalPotential, dimensions::Mass, normalization::UnitMetric>({1.0}, scales);
        const std::array<double, 1> state{0.3 * scales.specific_energy.value()}, residual{0.7 * scales.mass.value()};
        CHECK(factors.state_norm_squared(state) == Catch::Approx(0.09).epsilon(2.0e-13));
        CHECK(factors.residual_norm_squared(residual) == Catch::Approx(0.49).epsilon(2.0e-13));
    }
}