#include <algorithm>
#include <cmath>
#include <cstring>
#include <limits>
#include <numbers>
#include <stdexcept>

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

import mean_field;
import test_helpers;

namespace spherical_harmonic_surface_tests {
    namespace deformation = mean_field::deformation;
    namespace field       = mean_field::field;
    namespace domain      = mean_field::utils::domain;
    using Schema          = domain::CoreEnvelopeVacuumDomainSchema;

    int ranks() {
        int count = 0;
        MPI_Comm_size(MPI_COMM_WORLD, &count);
        return count;
    }

    auto context(const mean_field::fem::FEM &fem) {
        return deformation::SurfaceDeformationCompilationContext{
            *fem.surfaceDeformationFes, field::make_stellar_surface_scalar_dof_map<Schema>(*fem.surfaceDeformationFes)
        };
    }

    int modeIndex(
        const deformation::PreparedSphericalHarmonicSurface &surface,
        int degree,
        int order
    ) {
        int index = 0;
        for (const auto mode : surface.modes()) {
            if (mode.degree == degree && mode.order == order)
                return index;
            ++index;
        }
        throw std::logic_error("Required analytic harmonic mode was not compiled.");
    }

    double relativeDifference(
        const mfem::Vector &a,
        const mfem::Vector &b
    ) {
        REQUIRE(a.Size() == b.Size());
        mfem::Vector difference(a);
        difference -= b;
        return difference.Norml2() / std::max(b.Norml2(), 1e-300);
    }

    void fill(
        mfem::Vector &values,
        double scale = 1.
    ) {
        for (int i = 0; i < values.Size(); ++i)
            values(i) = scale * std::sin(.37 * (i + 1));
    }
} // namespace spherical_harmonic_surface_tests

TEST_CASE(
    "Spherical Harmonic Surface Validates Physical Mode Selection And Type Contracts",
    tags::spherical_harmonic_surface_validation
) {
    namespace d = mean_field::deformation;
    STATIC_CHECK(d::SurfaceDeformationPrescription<d::SphericalHarmonicSurface>);
    STATIC_CHECK(d::PreparedSurfaceDeformationPrescription<d::PreparedSphericalHarmonicSurface>);
    STATIC_CHECK(d::SurfaceDeformationCompilable<d::SphericalHarmonicSurface, d::SurfaceDeformationCompilationContext>);
    STATIC_CHECK_FALSE(d::PreparedSphericalHarmonicSurface::hasNodalRadialParameters);

    const d::SphericalHarmonicSurface full;
    const d::SphericalHarmonicSurface axisymmetric(
        d::SphericalHarmonicSurfaceOptions{.maximumDegree = 4, .maximumAzimuthalOrder = 0}
    );
    const d::SphericalHarmonicSurface orderOne(
        d::SphericalHarmonicSurfaceOptions{.maximumDegree = 4, .maximumAzimuthalOrder = 1}
    );
    const d::SphericalHarmonicSurface orderTwo(
        d::SphericalHarmonicSurfaceOptions{.maximumDegree = 4, .maximumAzimuthalOrder = 2}
    );
    const d::SphericalHarmonicSurface monopole(d::SphericalHarmonicSurfaceOptions{.maximumDegree = 0});
    CHECK(full.modeCount() == 25);
    CHECK(axisymmetric.modeCount() == 5);
    CHECK(orderOne.modeCount() == 13);
    CHECK(orderTwo.modeCount() == 19);
    CHECK(monopole.modeCount() == 1);
    CHECK(full.referenceCenter().Size() == 3);
    CHECK(full.referenceCenter().Norml2() == 0.);
    CHECK(full.descriptor().motionKind == d::SurfaceMotionKind::Radial);
    CHECK(full.descriptor().linearOnReferenceGeometry);
    CHECK(full.descriptor().supportsExactNewtonLinearization());

    for (const auto options :
         {d::SphericalHarmonicSurfaceOptions{.maximumDegree = -1},
          d::SphericalHarmonicSurfaceOptions{.maximumDegree = 4, .maximumAzimuthalOrder = -1},
          d::SphericalHarmonicSurfaceOptions{.maximumDegree = 4, .maximumAzimuthalOrder = 5},
          d::SphericalHarmonicSurfaceOptions{.maximumDegree = std::numeric_limits<int>::max()}})
        CHECK_THROWS_AS(d::SphericalHarmonicSurface(options), std::invalid_argument);
    mfem::Vector center(2);
    center = 0.;
    CHECK_THROWS_AS(d::SphericalHarmonicSurface(center), std::invalid_argument);
    center.SetSize(3);
    center    = 0.;
    center(1) = std::numeric_limits<double>::quiet_NaN();
    CHECK_THROWS_AS(d::SphericalHarmonicSurface(center), std::invalid_argument);
}

TEST_CASE(
    "Spherical Harmonic Surface Matches Analytic Real Harmonics And Radial Lengths",
    tags::spherical_harmonic_surface_validation
) {
    namespace t = spherical_harmonic_surface_tests;
    namespace d = mean_field::deformation;
    auto args   = test_utils::setup_args();
    auto fem    = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(fem.okay());
    const auto context = t::context(fem);
    if (t::ranks() != 1) {
        CHECK_THROWS_AS(
            d::compileSurfaceDeformationPrescription(d::SphericalHarmonicSurface{}, context), std::invalid_argument
        );
        return;
    }
    const auto surface = d::compileSurfaceDeformationPrescription(d::SphericalHarmonicSurface{}, context);
    CHECK(surface.parameterCount() == 25);
    CHECK(surface.globalParameterCount() == 25);
    CHECK(surface.globalParameterOffset() == 0);
    CHECK(surface.surfaceDisplacementSize() == 3 * surface.surfaceNodeCount());
    REQUIRE(surface.surfaceNodeCount() > surface.parameterCount());
    CHECK(surface.nodalBasis().Height() == surface.surfaceNodeCount());
    CHECK(surface.nodalBasis().Width() == surface.parameterCount());
    const double pi = std::numbers::pi;
    for (int node = 0; node < surface.surfaceNodeCount(); ++node) {
        const double x = surface.radialDirection(node, 0), y = surface.radialDirection(node, 1),
                     z   = surface.radialDirection(node, 2);
        const auto check = [&](int l, int m, double exact) {
            CHECK(std::abs(surface.nodalBasis()(node, t::modeIndex(surface, l, m)) - exact) < 2e-14);
        };
        check(0, 0, 1. / std::sqrt(4. * pi));
        check(1, 0, std::sqrt(3. / (4. * pi)) * z);
        check(1, 1, -std::sqrt(3. / (4. * pi)) * x);
        check(1, -1, -std::sqrt(3. / (4. * pi)) * y);
        check(2, 0, std::sqrt(5. / (16. * pi)) * (3. * z * z - 1.));
        check(2, -2, std::sqrt(15. / (4. * pi)) * x * y);
        check(2, 2, std::sqrt(15. / (16. * pi)) * (x * x - y * y));
    }
    mfem::Vector parameters(surface.parameterCount()), displacement(surface.surfaceDisplacementSize());
    parameters    = 0.;
    parameters(0) = .125 * std::sqrt(4. * pi);
    surface.buildSurfaceDisplacement(parameters, displacement);
    for (int node = 0; node < surface.surfaceNodeCount(); ++node)
        for (int component = 0; component < 3; ++component)
            CHECK(
                std::abs(
                    displacement(surface.surfaceDisplacementDof(node, component)) -
                    .125 * surface.radialDirection(node, component)
                ) < 1e-15
            );

    const auto axisymmetric = d::compileSurfaceDeformationPrescription(
        d::SphericalHarmonicSurface(d::SphericalHarmonicSurfaceOptions{.maximumDegree = 4, .maximumAzimuthalOrder = 0}),
        context
    );
    const auto restricted = d::compileSurfaceDeformationPrescription(
        d::SphericalHarmonicSurface(d::SphericalHarmonicSurfaceOptions{.maximumDegree = 4, .maximumAzimuthalOrder = 1}),
        context
    );
    REQUIRE(axisymmetric.parameterCount() == 5);
    REQUIRE(restricted.parameterCount() == 13);
    for (int j = 0; j < axisymmetric.parameterCount(); ++j) {
        CHECK(axisymmetric.modes()[j].degree == j);
        CHECK(axisymmetric.modes()[j].order == 0);
        for (int i = 0; i < surface.surfaceNodeCount(); ++i)
            CHECK(axisymmetric.nodalBasis()(i, j) == surface.nodalBasis()(i, t::modeIndex(surface, j, 0)));
    }
    for (const auto mode : restricted.modes())
        CHECK(std::abs(mode.order) <= 1);
    const d::SphericalHarmonicSurface unresolved(
        d::SphericalHarmonicSurfaceOptions{.maximumDegree = surface.surfaceNodeCount()}
    );
    CHECK_THROWS_AS(d::compileSurfaceDeformationPrescription(unresolved, context), std::invalid_argument);
}

TEST_CASE(
    "Spherical Harmonic Surface Preserves Its Exact Derivative And Dual Work",
    tags::spherical_harmonic_surface_validation
) {
    namespace t = spherical_harmonic_surface_tests;
    namespace d = mean_field::deformation;
    auto args   = test_utils::setup_args();
    auto fem    = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(fem.okay());
    const auto context = t::context(fem);
    if (t::ranks() != 1) {
        CHECK_THROWS_AS(
            d::compileSurfaceDeformationPrescription(d::SphericalHarmonicSurface{}, context), std::invalid_argument
        );
        return;
    }
    const auto surface = d::compileSurfaceDeformationPrescription(d::SphericalHarmonicSurface{}, context);
    const mfem::DenseMatrix originalBasis(surface.nodalBasis());
    mfem::Vector parameters(surface.parameterCount()), direction(surface.parameterCount()),
        dual(surface.surfaceDisplacementSize()), action(surface.surfaceDisplacementSize()),
        pulledBack(surface.parameterCount()), second(surface.parameterCount());
    t::fill(parameters, .03);
    t::fill(direction, .07);
    t::fill(dual);
    surface.applyJacobian(parameters, direction, action);
    surface.applyJacobianTranspose(parameters, dual, pulledBack);
    const double scale = std::max({1., action.Norml2() * dual.Norml2(), direction.Norml2() * pulledBack.Norml2()});
    CHECK(std::abs((action * dual) - (direction * pulledBack)) < 2e-13 * scale);

    constexpr double epsilon = 1e-5;
    mfem::Vector plus(parameters), minus(parameters), plusValue(action.Size()), minusValue(action.Size());
    plus.Add(epsilon, direction);
    minus.Add(-epsilon, direction);
    surface.buildSurfaceDisplacement(plus, plusValue);
    surface.buildSurfaceDisplacement(minus, minusValue);
    plusValue -= minusValue;
    plusValue /= 2. * epsilon;
    CHECK(t::relativeDifference(plusValue, action) < 2e-10);
    second = 7.;
    surface.applyPullbackDerivative(parameters, direction, dual, second);
    CHECK(second.Norml2() == 0.);
    CHECK(
        std::memcmp(
            originalBasis.GetData(), surface.nodalBasis().GetData(),
            sizeof(mfem::real_t) * originalBasis.Height() * originalBasis.Width()
        ) == 0
    );

    mfem::Vector wrong(surface.parameterCount() + 1);
    wrong = 0.;
    CHECK_THROWS_AS(surface.buildSurfaceDisplacement(wrong, action), std::invalid_argument);
    CHECK_THROWS_AS(surface.applyJacobian(parameters, wrong, action), std::invalid_argument);
    CHECK_THROWS_AS(surface.applyJacobianTranspose(parameters, dual, wrong), std::invalid_argument);
    CHECK_THROWS_AS(surface.applyPullbackDerivative(parameters, direction, dual, wrong), std::invalid_argument);
    direction(0) = std::numeric_limits<double>::quiet_NaN();
    CHECK_THROWS_AS(surface.applyJacobian(parameters, direction, action), std::invalid_argument);
    dual(0) = std::numeric_limits<double>::infinity();
    CHECK_THROWS_AS(surface.applyJacobianTranspose(parameters, dual, pulledBack), std::invalid_argument);
    // Finite entries can still overflow an accumulated modal force; accepting
    // an infinite dual would contaminate the complete Newton merit silently.
    for (int node = 0; node < surface.surfaceNodeCount(); ++node)
        for (int component = 0; component < 3; ++component)
            dual(surface.surfaceDisplacementDof(node, component)) =
                std::numeric_limits<double>::max() * surface.radialDirection(node, component);
    for (int i = 0; i < dual.Size(); ++i)
        REQUIRE(std::isfinite(dual(i)));
    CHECK_THROWS_AS(surface.applyJacobianTranspose(parameters, dual, pulledBack), std::invalid_argument);
}

TEST_CASE(
    "Spherical Harmonic Surface Rejects Aliased Modes Even With Enough Nodes",
    tags::spherical_harmonic_surface_validation
) {
    namespace t = spherical_harmonic_surface_tests;
    namespace d = mean_field::deformation;
    if (t::ranks() != 1) {
        SUCCEED("MPI compilation rejection is covered by the full surface fixture.");
        return;
    }
    auto serial = mfem::Mesh::MakeCartesian3D(1, 1, 1, mfem::Element::HEXAHEDRON, 1., 1., 1.);
    for (int b = 0; b < serial.GetNBE(); ++b)
        serial.GetBdrElement(b)->SetAttribute(1);
    mfem::ParMesh mesh(MPI_COMM_WORLD, serial);
    mfem::H1_FECollection collection(1, 3);
    mfem::ParFiniteElementSpace space(&mesh, &collection);
    const auto map = t::field::make_stellar_surface_scalar_dof_map<t::Schema>(space);
    const d::SurfaceDeformationCompilationContext context(space, map);
    mfem::Vector center(3);
    center = .5;
    const d::SphericalHarmonicSurface axisymmetric(
        center, d::SphericalHarmonicSurfaceOptions{.maximumDegree = 2, .maximumAzimuthalOrder = 0}
    );
    // Eight cube corners offer only two distinct cos(theta), so the three
    // axisymmetric polynomials cannot define three independent coordinates.
    if (t::ranks() == 1)
        REQUIRE(map.local_size() == 8);
    CHECK_THROWS_AS(d::compileSurfaceDeformationPrescription(axisymmetric, context), std::invalid_argument);
}

TEST_CASE(
    "Spherical Harmonic Domain Couples All Volume Fields Through The Existing Extensions",
    tags::spherical_harmonic_surface_validation
) {
    namespace t = spherical_harmonic_surface_tests;
    namespace d = mean_field::deformation;
    auto args   = test_utils::setup_args();
    auto fem    = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(fem.okay());
    const auto make = [&]() {
        return d::compileDomainDeformation(
            d::SphericalHarmonicSurface(d::SphericalHarmonicSurfaceOptions{.maximumDegree = 2}),
            d::PowerLawRadialInteriorExtension{}, d::FixedInfinityRadialVacuumExtension{}, fem
        );
    };
    if (t::ranks() != 1) {
        CHECK_THROWS_AS(make(), std::invalid_argument);
        return;
    }
    auto harmonic       = make();
    const auto &surface = harmonic.surfaceDeformationPrescription();
    mfem::Vector center(3);
    center     = 0.;
    auto nodal = d::compileDomainDeformation(
        d::NodalRadialSurface(center), d::PowerLawRadialInteriorExtension{}, d::FixedInfinityRadialVacuumExtension{},
        fem
    );
    mfem::Vector parameters(harmonic.parameterCount()), direction(harmonic.parameterCount()),
        nodalParameters(nodal.parameterCount()), volume(harmonic.volumeDisplacementSize()),
        nodalVolume(nodal.volumeDisplacementSize()), action(volume.Size()), dual(volume.Size()),
        pulledBack(parameters.Size()), pullbackDerivative(parameters.Size());
    REQUIRE(harmonic.parameterCount() == 9);
    REQUIRE(harmonic.surfaceDisplacementSize() == nodal.surfaceDisplacementSize());
    REQUIRE(harmonic.volumeDisplacementSize() == nodal.volumeDisplacementSize());
    t::fill(parameters, .001);
    t::fill(direction, .004);
    t::fill(dual);
    surface.nodalBasis().Mult(parameters, nodalParameters);
    harmonic.buildVolumeDisplacement(parameters, volume);
    nodal.buildVolumeDisplacement(nodalParameters, nodalVolume);
    CHECK(t::relativeDifference(volume, nodalVolume) < 3e-13);
    harmonic.applyJacobian(parameters, direction, action);
    harmonic.applyJacobianTranspose(parameters, dual, pulledBack);
    CHECK(
        std::abs((action * dual) - (direction * pulledBack)) <
        1e-11 * std::max({1., action.Norml2() * dual.Norml2(), direction.Norml2() * pulledBack.Norml2()})
    );
    harmonic.applyPullbackDerivative(parameters, direction, dual, pullbackDerivative);
    CHECK(pullbackDerivative.Norml2() == 0.);

    constexpr double epsilon = 1e-5;
    mfem::Vector plus(parameters), minus(parameters), plusVolume(volume.Size()), minusVolume(volume.Size());
    plus.Add(epsilon, direction);
    minus.Add(-epsilon, direction);
    harmonic.buildVolumeDisplacement(plus, plusVolume);
    harmonic.buildVolumeDisplacement(minus, minusVolume);
    plusVolume -= minusVolume;
    plusVolume /= 2. * epsilon;
    CHECK(t::relativeDifference(plusVolume, action) < 2e-9);
    const auto infinity =
        t::field::make_scalar_boundary_dof_map<t::domain::InfinitySurface, t::Schema>(*fem.surfaceDeformationFes);
    for (int i = 0; i < infinity.local_size(); ++i)
        for (int component = 0; component < 3; ++component)
            CHECK(volume(infinity.volume_true_dof(i) + component * fem.surfaceDeformationFes->GetTrueVSize()) == 0.);
}