#include <cmath>
#include <concepts>
#include <numbers>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>
#include <mpi.h>

import mean_field;
import test_helpers;

namespace harmonic_surface_metric_tests {
    // Use analytic Cartesian formulas and direct boundary quadrature, rather
    // than either the harmonic basis implementation or its assembled mass
    // action. This distinguishes vector trace work from scalar radial work.
    double analyticReferenceWork(
        const mean_field::fem::FEM &fem,
        int degree
    ) {
        mfem::VectorFunctionCoefficient coefficient(3, [degree](const mfem::Vector &position, mfem::Vector &value) {
            const double radius = position.Norml2();
            if (!(radius > 0.)) {
                // Interior center values cannot affect the boundary trace.
                value.SetSize(3);
                value = 0.;
                return;
            }
            const double z = position(2) / radius;
            const double harmonic =
                degree == 0 ? 1. / std::sqrt(4. * std::numbers::pi) : std::sqrt(3. / (4. * std::numbers::pi)) * z;
            value.SetSize(3);
            for (int component = 0; component < 3; ++component)
                value(component) = harmonic * position(component) / radius;
        });
        mfem::ParGridFunction tangent(fem.displacementFes.get());
        tangent.ProjectCoefficient(coefficient);
        mfem::Array<int> dofs;
        mfem::Vector local, shape;
        long double work = 0.;
        for (int b = 0; b < fem.mesh->GetNBE(); ++b) {
            if (fem.mesh->GetBdrAttribute(b) != 1)
                continue;
            const auto &element  = *fem.displacementFes->GetBE(b);
            auto &transformation = *fem.displacementFes->GetBdrElementTransformation(b);
            REQUIRE(fem.displacementFes->GetBdrElementVDofs(b, dofs) == nullptr);
            tangent.GetSubVector(dofs, local);
            REQUIRE(local.Size() == 3 * element.GetDof());
            shape.SetSize(element.GetDof());
            // Match the declared FE quadrature rule, not the assembled matrix
            // or its diagonal; the integrand is independently reconstructed.
            const auto &rule = mfem::MassIntegrator::GetRule(element, element, transformation);
            for (int q = 0; q < rule.GetNPoints(); ++q) {
                const auto &point = rule.IntPoint(q);
                transformation.SetIntPoint(&point);
                element.CalcShape(point, shape);
                long double normSquared = 0.;
                for (int component = 0; component < 3; ++component) {
                    long double value = 0.;
                    for (int i = 0; i < element.GetDof(); ++i)
                        value += static_cast<long double>(shape(i)) * local(i + component * element.GetDof());
                    normSquared += value * value;
                }
                work += point.weight * transformation.Weight() * normSquared;
            }
        }
        REQUIRE(std::isfinite(work));
        REQUIRE(work > 0.);
        return static_cast<double>(work);
    }

    // A valid surface implementation does not automatically promise the
    // canonical nodal ordering required by this particular Riesz metric.
    class SurfaceWithoutTraceMetadata final {
    public:
        explicit SurfaceWithoutTraceMetadata(mean_field::deformation::PreparedSphericalHarmonicSurface surface)
            : m_surface(std::move(surface)) {
        }
        auto descriptor() const noexcept {
            return m_surface.descriptor();
        }
        int parameterCount() const noexcept {
            return m_surface.parameterCount();
        }
        int surfaceDisplacementSize() const noexcept {
            return m_surface.surfaceDisplacementSize();
        }
        void buildSurfaceDisplacement(
            const mfem::Vector &x,
            mfem::Vector &u
        ) const {
            m_surface.buildSurfaceDisplacement(x, u);
        }
        void applyJacobian(
            const mfem::Vector &x,
            const mfem::Vector &v,
            mfem::Vector &u
        ) const {
            m_surface.applyJacobian(x, v, u);
        }
        void applyJacobianTranspose(
            const mfem::Vector &x,
            const mfem::Vector &u,
            mfem::Vector &v
        ) const {
            m_surface.applyJacobianTranspose(x, u, v);
        }
        void applyPullbackDerivative(
            const mfem::Vector &x,
            const mfem::Vector &v,
            const mfem::Vector &u,
            mfem::Vector &result
        ) const {
            m_surface.applyPullbackDerivative(x, v, u, result);
        }

    private:
        mean_field::deformation::PreparedSphericalHarmonicSurface m_surface;
    };
    static_assert(mean_field::deformation::PreparedSurfaceDeformationPrescription<SurfaceWithoutTraceMetadata>);
} // namespace harmonic_surface_metric_tests

TEST_CASE(
    "Harmonic Surface Coordinates Receive Their Own Primal And Dual Riesz Scales Before Physical Preparation",
    "[normalization][spherical_harmonic][integration][serial]"
) {
    using namespace mean_field;
    int ranks = 0;
    MPI_Comm_size(MPI_COMM_WORLD, &ranks);
    if (ranks != 1)
        SKIP("The current radial domain extension supports one MPI rank.");

    // Both configurations have the same ambient trace space. Only the
    // independent surface coordinate count changes, so a nodal-diagonal
    // shortcut would fail the block-size and scale checks below.
    deformation::SphericalHarmonicSurface prescription({.maximumDegree = 4});
    int expectedModes = 25;
    SECTION("All real azimuthal modes") {
    }
    SECTION("Axisymmetric modes") {
        prescription  = deformation::SphericalHarmonicSurface({.maximumDegree = 4, .maximumAzimuthalOrder = 0});
        expectedModes = 5;
    }

    const auto args     = test_utils::setup_args();
    auto finiteElements = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(finiteElements.okay());
    const normalization::PhysicalRieszDiagonal policy{dimensions::LengthValue{utils::RADIUS}, utils::G};
    auto stellarModel = model::StellarModel(
        eos::Polytrope({.n = 1.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}})
    );
    auto problem = equilibrium::discretize(
        stellarModel, equilibrium::makeStellarDiscretization(std::move(finiteElements), policy, prescription)
    );
    using Problem               = std::remove_cvref_t<decltype(problem)>;
    using Form                  = typename Problem::FormType;
    using ShapeValue            = utils::blocks::surface_deformation::parameters::value;
    using ShapeResidual         = utils::blocks::surface_deformation::shape_equilibrium::residual;
    constexpr int valueBlock    = utils::blocks::type_index_v<ShapeValue, typename Form::value_blocks>;
    constexpr int residualBlock = utils::blocks::type_index_v<ShapeResidual, typename Form::residual_blocks>;
    STATIC_CHECK(normalization::NormalizableStellarEquilibriumProblem<Problem>);
    const auto &layout       = problem.GetManifest().layout();
    const int valueOffset    = layout.value_offsets()[valueBlock];
    const int residualOffset = layout.residual_offsets()[residualBlock];
    REQUIRE(layout.value_offsets()[valueBlock + 1] - valueOffset == expectedModes);
    REQUIRE(layout.residual_offsets()[residualBlock + 1] - residualOffset == expectedModes);
    REQUIRE(problem.StateSize() == problem.EquationSize());

    const auto &physical = problem.GetPhysicalOperator();
    const auto &domain   = physical.GetDomainDeformation();
    REQUIRE_FALSE(domain.hasNodalRadialParameters());
    REQUIRE(domain.parameterCount() == expectedModes);
    REQUIRE_FALSE(physical.IsPrepared());
    const auto generation = problem.GetPreparationGeneration();
    auto normalized       = normalization::makeNormalizedStellarEquilibriumOperator(problem);
    CHECK_FALSE(normalized.IsPrepared());
    CHECK_FALSE(physical.IsPrepared());
    CHECK(problem.GetPreparationGeneration() == generation);

    const auto diagonal =
        deformation::assembleSurfaceParameterMassDiagonal(physical.GetHydrostaticOperator().GetFEM(), domain);
    REQUIRE(diagonal.Size() == expectedModes);
    const auto &referenceFEM = physical.GetHydrostaticOperator().GetFEM();
    CHECK(
        diagonal(0) ==
        Catch::Approx(harmonic_surface_metric_tests::analyticReferenceWork(referenceFEM, 0)).epsilon(5e-13)
    );
    const int axialDipole = expectedModes == 5 ? 1 : 2;
    CHECK(
        diagonal(axialDipole) ==
        Catch::Approx(harmonic_surface_metric_tests::analyticReferenceWork(referenceFEM, 1)).epsilon(5e-13)
    );
    const auto scales           = normalization::deriveStellarCharacteristicScales(policy, stellarModel);
    const double valueScale     = normalization::physicalScale<ShapeValue>(scales);
    const double residualScale  = normalization::physicalScale<ShapeResidual>(scales);
    const auto &stateFactors    = normalized.GetNormalization().StateFactors();
    const auto &residualFactors = normalized.GetNormalization().ResidualFactors();
    REQUIRE(stateFactors.Size() == problem.StateSize());
    REQUIRE(residualFactors.Size() == problem.EquationSize());
    for (int i = 0; i < expectedModes; ++i) {
        REQUIRE(std::isfinite(diagonal[i]));
        REQUIRE(diagonal[i] > 0.0);
        CHECK(stateFactors[valueOffset + i] == Catch::Approx(std::sqrt(diagonal[i]) / valueScale).epsilon(3.e-14));
        CHECK(
            residualFactors[residualOffset + i] ==
            Catch::Approx(1.0 / (residualScale * std::sqrt(diagonal[i]))).epsilon(3.e-14)
        );
        // Primal and dual factors act on different root block offsets. Their
        // Gram weights cancel in the work pairing, leaving physical scales.
        CHECK(
            stateFactors[valueOffset + i] * residualFactors[residualOffset + i] ==
            Catch::Approx(1.0 / (valueScale * residualScale)).epsilon(3.e-14)
        );
    }
    bool allFactorsValid = true;
    for (int i = 0; i < stateFactors.Size(); ++i)
        allFactorsValid = allFactorsValid && std::isfinite(stateFactors[i]) && stateFactors[i] > 0.0 &&
                          std::isfinite(residualFactors[i]) && residualFactors[i] > 0.0;
    CHECK(allFactorsValid);
    CHECK_FALSE(physical.IsPrepared());
    CHECK(problem.GetPreparationGeneration() == generation);
}

TEST_CASE(
    "Harmonic Reference Metric Rejects Unrelated Current Spaces And Unadvertised Trace Layouts",
    "[normalization][spherical_harmonic][integration][serial]"
) {
    using namespace mean_field;
    int ranks = 0;
    MPI_Comm_size(MPI_COMM_WORLD, &ranks);
    if (ranks != 1)
        SKIP("The current radial domain extension supports one MPI rank.");
    const auto args     = test_utils::setup_args();
    auto finiteElements = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(finiteElements.okay());
    deformation::SphericalHarmonicSurface prescription({.maximumDegree = 2});
    auto prepared = deformation::compileDomainDeformation(
        prescription, deformation::PowerLawRadialInteriorExtension{}, deformation::FixedInfinityRadialVacuumExtension{},
        finiteElements
    );
    deformation::PreparedDomainDeformationRuntime runtime(std::move(prepared));
    REQUIRE(runtime.matchesCurrentDiscretization());
    REQUIRE(runtime.supportsSurfaceParameterMetric());

    SECTION("A different current FEM with identical dimensions is not interchangeable") {
        auto other = fem::setup_fem(args.mesh_file, args, 0);
        REQUIRE(other.okay());
        REQUIRE(other.surfaceDeformationFes->GetTrueVSize() == finiteElements.surfaceDeformationFes->GetTrueVSize());
        REQUIRE(other.displacementFes->GetTrueVSize() == finiteElements.displacementFes->GetTrueVSize());
        REQUIRE(other.mesh.get() != finiteElements.mesh.get());
        CHECK_THROWS_AS(deformation::assembleSurfaceParameterMassDiagonal(other, runtime), std::invalid_argument);
        CHECK(runtime.matchesCurrentDiscretization());
    }

    SECTION("A surface tangent alone cannot certify canonical nodal vector ordering") {
        using Schema = utils::domain::CoreEnvelopeVacuumDomainSchema;
        const deformation::SurfaceDeformationCompilationContext surfaceContext{
            *finiteElements.surfaceDeformationFes,
            field::make_stellar_surface_scalar_dof_map<Schema>(*finiteElements.surfaceDeformationFes)
        };
        auto surface          = deformation::compileSurfaceDeformationPrescription(prescription, surfaceContext);
        auto extensionContext = deformation::makeRadialDeformationExtensionCompilationContext<Schema>(
            *finiteElements.surfaceDeformationFes, *finiteElements.displacementFes, *finiteElements.logicalReferenceMesh
        );
        auto custom = deformation::composePreparedDomainDeformation(
            harmonic_surface_metric_tests::SurfaceWithoutTraceMetadata(std::move(surface)),
            deformation::compileInteriorDeformationExtension(
                deformation::PowerLawRadialInteriorExtension{}, extensionContext
            ),
            deformation::compileVacuumDeformationExtension(
                deformation::FixedInfinityRadialVacuumExtension{}, extensionContext
            ),
            *finiteElements.surfaceDeformationFes, *finiteElements.displacementFes, *finiteElements.logicalReferenceMesh
        );
        STATIC_CHECK(deformation::PreparedDomainDeformationOperator<decltype(custom)>);
        REQUIRE_FALSE(custom.supportsSurfaceParameterMetric());
        deformation::PreparedDomainDeformationRuntime customRuntime(std::move(custom));
        REQUIRE(customRuntime.matchesCurrentDiscretization());
        REQUIRE_FALSE(customRuntime.supportsSurfaceParameterMetric());
        mfem::Vector parameters(customRuntime.parameterCount()), direction(parameters.Size()),
            trace(customRuntime.surfaceDisplacementSize());
        parameters = 0.;
        direction  = 1.;
        CHECK_NOTHROW(customRuntime.applySurfaceJacobian(parameters, direction, trace));
        CHECK(trace.Norml2() > 0.);
        CHECK_THROWS_AS(
            deformation::assembleSurfaceParameterMassDiagonal(finiteElements, customRuntime), std::invalid_argument
        );
    }
}