#include <algorithm>
#include <cmath>
#include <cstdint>
#include <stdexcept>
#include <type_traits>
#include <utility>

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

import mean_field;
import test_helpers;

namespace spherical_harmonic_preconditioner_test {
    using namespace mean_field;

    [[nodiscard]] operators::StellarEquilibriumDependencies dependencies() {
        return {
            .discretization     = {.identity = 9201, .revision = 1},
            .density            = {.identity = 9203, .revision = 1},
            .surfaceDeformation = {.identity = 9209, .revision = 1},
            .gravityGradient    = {.identity = 9221, .revision = 1},
            .gravityPotential   = {.identity = 9227, .revision = 1},
            .enthalpy           = {.identity = 9239, .revision = 1},
            .bernoulliConstant  = {.identity = 9241, .revision = 1},
            .rotation           = {.identity = 9257, .revision = 1},
            .targetMass         = {.identity = 9277, .revision = 1}
        };
    }

    void checkClose(
        const mfem::Vector &actual,
        const mfem::Vector &expected
    ) {
        REQUIRE(actual.Size() == expected.Size());
        mfem::Vector difference(actual);
        difference -= expected;
        const double error = difference.Norml2();
        const double scale = std::max(actual.Norml2(), expected.Norml2());
        REQUIRE(std::isfinite(error));
        REQUIRE(std::isfinite(scale));
        INFO("inverse error = " << error << ", scale = " << scale);
        CHECK(error <= 2.e-12 + 3.e-11 * scale);
    }

    // Construct the advertised triangular surrogate's RHS from a known
    // correction. This checks the inverse's parameter ordering and material
    // coupling, without asking an approximate inverse to invert the full J.
    template <typename Prepared>
    void checkTriangularInverse(
        const Prepared &inverse,
        const int modes
    ) {
        const auto &couplings = inverse.GetCoupledOperator();
        const auto &offsets   = couplings.GetOffsets();
        REQUIRE(offsets[2] - offsets[1] == modes);
        REQUIRE(inverse.Width() == offsets[3]);
        REQUIRE(inverse.Height() == inverse.Width());
        REQUIRE(inverse.GetSurfaceDiagonal().Size() == modes);
        mfem::Vector expected(inverse.Height()), rhs(inverse.Width()), actual(inverse.Height());
        for (int i = 0; i < expected.Size(); ++i)
            expected(i) = 0.02 * std::sin(0.31 * (i + 1));
        const mfem::Vector density(expected.GetData(), offsets[1]);
        const mfem::Vector shape(expected.GetData() + offsets[1], modes);
        const mfem::Vector enthalpy(expected.GetData() + offsets[2], offsets[3] - offsets[2]);
        mfem::Vector densityCoupling(density.Size()), enthalpyCoupling(enthalpy.Size()), cross(density.Size());
        couplings.ApplySurfaceToMaterial(shape, densityCoupling, enthalpyCoupling);
        couplings.ApplyEnthalpyToDensity(enthalpy, cross);
        for (int i = 0; i < density.Size(); ++i)
            rhs(i) = inverse.GetDensityDiagonal()(i) * density(i) + densityCoupling(i) + cross(i);
        for (int i = 0; i < modes; ++i)
            rhs(offsets[1] + i) = inverse.GetSurfaceDiagonal()(i) * shape(i);
        for (int i = 0; i < enthalpy.Size(); ++i)
            rhs(offsets[2] + i) = inverse.GetEnthalpyDiagonal()(i) * enthalpy(i) + enthalpyCoupling(i);
        actual = 0.0;
        inverse.Mult(rhs, actual);
        checkClose(actual, expected);
        const mfem::Vector actualShape(actual.GetData() + offsets[1], modes);
        checkClose(actualShape, shape);
        rhs    = 0.0;
        actual = 1.0;
        inverse.Mult(rhs, actual);
        CHECK(actual.Norml2() == 0.0);
        CHECK(inverse.IsCurrent());
    }
} // namespace spherical_harmonic_preconditioner_test

TEST_CASE(
    "Harmonic Surface Mass And Calibrated Inverses Use Native Parameters And Reject Nodal H1",
    "[preconditioning][material_surface][spherical_harmonic][integration][serial]"
) {
    using namespace mean_field;
    using namespace spherical_harmonic_preconditioner_test;
    namespace backend = preconditioning::backend;
    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());
    const auto stellarModel = model::StellarModel(
        eos::Polytrope({.n = 1.5, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}})
    );
    // A restricted azimuthal set also rules out using (L+1)^2 as the block
    // length when the actual mode catalog has fewer independent coefficients.
    const deformation::SphericalHarmonicSurface prescription{{.maximumDegree = 2, .maximumAzimuthalOrder = 1}};
    constexpr int modes = 7;
    auto problem        = equilibrium::discretize(
        stellarModel,
        equilibrium::makeStellarDiscretization(std::move(finiteElements), normalization::Unnormalized{}, prescription)
    );
    mfem::Vector state(problem.StateSize());
    state                                                   = 0.0;
    const auto view                                         = problem.GetManifest().stateView(state);
    view.block(utils::blocks::density_field.mass_term)      = 1.0;
    view.block(utils::blocks::enthalpy_field.specific_term) = 0.7;
    view.block(utils::blocks::fixed_total_mass_constraint.mass_normalization_term) = 0.8;
    auto shape = view.block(utils::blocks::surface_deformation_field.parameters_term);
    REQUIRE(shape.Size() == modes);
    for (int i = 0; i < modes; ++i)
        shape(i) = 0.001 * std::sin(0.7 * (i + 1));
    mfem::Vector omega(3), center(3);
    omega                  = 0.0;
    center                 = 0.0;
    const auto preparation = problem.TryPrepare(state, dependencies(), physics::RigidRotation{omega, center});
    REQUIRE(preparation.has_value());
    const auto &physical = problem.GetPhysicalOperator();
    REQUIRE(physical.IsPrepared());
    REQUIRE_FALSE(physical.GetDomainDeformation().hasNodalRadialParameters());

    const auto block = preconditioning::materialSurfaceBlock(problem);
    STATIC_CHECK(
        std::is_same_v<
            typename std::remove_cvref_t<decltype(block)>::SurfaceSurrogate, preconditioning::SurfaceMassDiagonal>
    );
    auto prepared = preconditioning::prepare(problem, block);
    CHECK(prepared.GetStatistics().surfaceRieszAssemblies == 1);
    CHECK(prepared.GetStatistics().surfaceJacobianProbes == 0);
    CHECK(prepared.GetSurfaceDiagonalQuality().minimumAbsoluteEntryBeforeRegularization > 0.0);
    checkTriangularInverse(prepared, modes);
    const auto unchanged = prepared.Refresh(physical);
    CHECK_FALSE(unchanged.DidAnyWork());

    const auto calibratedBlock = preconditioning::materialSurfaceBlock(
        problem, backend::Diagonal{}, backend::Diagonal{}, preconditioning::SurfaceThenMaterialTriangular{},
        {.surfaceCalibration = {
             .target     = preconditioning::SurfaceRieszCalibrationTarget::surface_jacobian,
             .probeCount = 3,
             .objective  = preconditioning::SurfaceRieszCalibrationObjective::operator_action
         }}
    );
    auto calibrated         = preconditioning::prepare(problem, calibratedBlock);
    const auto &calibration = calibrated.GetSurfaceCalibration();
    REQUIRE(calibration.WasCalibrated());
    REQUIRE(std::isfinite(calibration.scale));
    REQUIRE(calibration.scale != 0.0);
    CHECK(calibration.probeCount == 3);
    CHECK(calibrated.GetStatistics().surfaceJacobianProbes == 3);
    for (int i = 0; i < modes; ++i)
        CHECK(
            calibrated.GetSurfaceDiagonal()(i) ==
            Catch::Approx(calibration.scale * prepared.GetSurfaceDiagonal()(i)).epsilon(2.e-13)
        );
    checkTriangularInverse(calibrated, modes);

    using FixedAMG = backend::HypreBoomerAMG<backend::FixedCycles>;
    const auto h1  = preconditioning::materialSurfaceBlock(
        problem, backend::Diagonal{}, FixedAMG{backend::FixedCycles{.cycles = 1}},
        preconditioning::ApproximateMaterialSurfaceLDU{}, preconditioning::SurfaceH1MassStiffness{}
    );
    CHECK_THROWS_WITH(
        preconditioning::prepare(problem, h1), Catch::Matchers::ContainsSubstring("requires nodal radial parameters")
    );
    CHECK(prepared.IsCurrent());
    CHECK(calibrated.IsCurrent());
}