#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <optional>
#include <span>
#include <string_view>
#include <type_traits>
#include <utility>

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

import mean_field;
import test_helpers;

namespace spherical_harmonic_equilibrium_test {
    using namespace mean_field;
    using Block = operators::RootBlockDescriptor;

    [[nodiscard]] const Block &findBlock(
        const std::span<const Block> blocks,
        const std::string_view name
    ) {
        const auto found =
            std::find_if(blocks.begin(), blocks.end(), [name](const Block &block) { return block.stableId == name; });
        REQUIRE(found != blocks.end());
        return *found;
    }

    [[nodiscard]] mfem::Vector readBlock(
        const mfem::Vector &vector,
        const Block &block
    ) {
        REQUIRE(block.offset >= 0);
        REQUIRE(block.offset + block.size <= vector.Size());
        mfem::Vector result(block.size);
        for (int i = 0; i < block.size; ++i)
            result(i) = vector(block.offset + i);
        return result;
    }

    void writeBlock(
        mfem::Vector &vector,
        const Block &block,
        const mfem::Vector &values
    ) {
        REQUIRE(values.Size() == block.size);
        for (int i = 0; i < block.size; ++i)
            vector(block.offset + i) = values(i);
    }

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

    void checkBlocks(
        const mfem::Vector &actual,
        const mfem::Vector &expected,
        const std::span<const Block> blocks,
        const double relative = 2.0e-9,
        const double absolute = 2.0e-12
    ) {
        for (const Block &block : blocks) {
            INFO("residual block = " << block.stableId);
            checkClose(readBlock(actual, block), readBlock(expected, block), relative, absolute);
        }
    }

    [[nodiscard]] operators::StellarEquilibriumDependencies dependencies(const std::uint64_t revision) {
        return {
            .discretization     = {.identity = 8101, .revision = 1},
            .density            = {.identity = 8103, .revision = revision},
            .surfaceDeformation = {.identity = 8107, .revision = revision},
            .gravityGradient    = {.identity = 8111, .revision = revision},
            .gravityPotential   = {.identity = 8117, .revision = revision},
            .enthalpy           = {.identity = 8123, .revision = revision},
            .bernoulliConstant  = {.identity = 8129, .revision = revision},
            .rotation           = {.identity = 8131, .revision = revision},
            .targetMass         = {.identity = 8147, .revision = 1}
        };
    }

    [[nodiscard]] physics::RigidRotation zeroRotation() {
        mfem::Vector omega(3), center(3);
        omega  = 0.0;
        center = 0.0;
        return {omega, center};
    }

    template <bool GeneratedRotation> [[nodiscard]] auto makeModel() {
        if constexpr (GeneratedRotation) {
            return model::StellarModel(
                eos::Polytrope({.n = 1.5, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
                integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
                integral::FixedAngularMomentum({.Jtotal = dimensions::AngularMomentumValue{0.01}})
            );
        } else {
            return model::StellarModel(
                eos::Polytrope({.n = 1.5, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
                integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}})
            );
        }
    }

    template <typename Problem>
    void prepare(
        Problem &problem,
        const mfem::Vector &state,
        const std::uint64_t revision
    ) {
        if constexpr (Problem::hasFixedAngularMomentum) {
            REQUIRE(problem.TryPrepare(state, dependencies(revision)).has_value());
        } else {
            REQUIRE(problem.TryPrepare(state, dependencies(revision), zeroRotation()).has_value());
        }
    }

    // The oracle is the existing nodal problem, with explicit coefficient and
    // dual-force maps. It never calls the native modal residual or Jacobian to
    // construct an expected coupled result.
    template <
        typename ModalProblem,
        typename NodalProblem>
    [[nodiscard]] mfem::Vector liftState(
        const ModalProblem &modal,
        const NodalProblem &nodal,
        const mfem::DenseMatrix &basis,
        const mfem::Vector &source
    ) {
        mfem::Vector result(nodal.StateSize());
        result = 0.0;
        for (const Block &destination : nodal.GetManifest().valueBlocks()) {
            const Block &origin       = findBlock(modal.GetManifest().valueBlocks(), destination.stableId);
            const mfem::Vector values = readBlock(source, origin);
            if (destination.stableId == "surface_deformation") {
                REQUIRE(basis.Width() == origin.size);
                REQUIRE(basis.Height() == destination.size);
                mfem::Vector lifted(destination.size);
                basis.Mult(values, lifted);
                writeBlock(result, destination, lifted);
            } else {
                writeBlock(result, destination, values);
            }
        }
        return result;
    }

    template <
        typename ModalProblem,
        typename NodalProblem>
    [[nodiscard]] mfem::Vector pullResidual(
        const ModalProblem &modal,
        const NodalProblem &nodal,
        const mfem::DenseMatrix &basis,
        const mfem::Vector &source
    ) {
        mfem::Vector result(modal.EquationSize());
        result = 0.0;
        for (const Block &destination : modal.GetManifest().residualBlocks()) {
            const Block &origin       = findBlock(nodal.GetManifest().residualBlocks(), destination.stableId);
            const mfem::Vector values = readBlock(source, origin);
            if (destination.stableId == "surface_shape_balance") {
                REQUIRE(basis.Height() == origin.size);
                REQUIRE(basis.Width() == destination.size);
                mfem::Vector pulled(destination.size);
                basis.MultTranspose(values, pulled);
                writeBlock(result, destination, pulled);
            } else {
                writeBlock(result, destination, values);
            }
        }
        return result;
    }

    template <bool GeneratedRotation> void checkCoupledEquivalence() {
        const utils::Args args = test_utils::setup_args();
        fem::FEM modalFem      = fem::setup_fem(args.mesh_file, args, 0);
        fem::FEM nodalFem      = fem::setup_fem(args.mesh_file, args, 0);
        REQUIRE(modalFem.okay());
        REQUIRE(nodalFem.okay());

        const deformation::SphericalHarmonicSurface prescription{{.maximumDegree = 4}};
        using Schema          = utils::domain::CoreEnvelopeVacuumDomainSchema;
        const auto surfaceMap = field::make_stellar_surface_scalar_dof_map<Schema>(*modalFem.surfaceDeformationFes);
        const auto preparedSurface = deformation::compileSurfaceDeformationPrescription(
            prescription, deformation::SurfaceDeformationCompilationContext{*modalFem.surfaceDeformationFes, surfaceMap}
        );
        const mfem::DenseMatrix basis(preparedSurface.nodalBasis());
        REQUIRE(basis.Width() == 25);

        const normalization::PhysicalRieszDiagonal policy{dimensions::LengthValue{1.0}};
        const auto stellarModel = makeModel<GeneratedRotation>();
        auto modal              = equilibrium::discretize(
            stellarModel, equilibrium::makeStellarDiscretization(std::move(modalFem), policy, prescription)
        );
        auto nodal =
            equilibrium::discretize(stellarModel, equilibrium::makeStellarDiscretization(std::move(nodalFem), policy));
        using ModalProblem = std::remove_cvref_t<decltype(modal)>;
        STATIC_CHECK_FALSE(ModalProblem::hasFixedCentralDensity);
        STATIC_CHECK(ModalProblem::hasFixedAngularMomentum == GeneratedRotation);
        STATIC_CHECK(ModalProblem::symbolicallySquare);
        REQUIRE(modal.StateSize() == modal.EquationSize());
        const Block &shape     = findBlock(modal.GetManifest().valueBlocks(), "surface_deformation");
        const Block &shapeRows = findBlock(modal.GetManifest().residualBlocks(), "surface_shape_balance");
        REQUIRE(shape.size == 25);
        REQUIRE(shapeRows.size == shape.size);
        CHECK(nodal.StateSize() - modal.StateSize() == basis.Height() - basis.Width());

        mfem::Vector state(modal.StateSize()), direction(modal.StateSize());
        state = 0.0;
        for (int i = 0; i < direction.Size(); ++i)
            direction(i) = 0.02 * std::sin(0.31 * (i + 1));
        const auto view                                         = modal.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;
        if constexpr (GeneratedRotation) {
            view.block(utils::blocks::fixed_angular_momentum_constraint.angular_velocity_term) = 0.04;
            modal.GetManifest().stateView(direction).block(
                utils::blocks::fixed_angular_momentum_constraint.angular_velocity_term
            ) = 0.015;
        }
        // Nonzero orders on both real branches make an accidental m=0-only
        // integration fail. Positive h keeps fractional powers differentiable;
        // its nonzero boundary value is handled by the physical replacement rows.
        for (int mode = 0; mode < shape.size; ++mode) {
            state(shape.offset + mode)     = 0.001 * std::sin(0.7 * (mode + 1));
            direction(shape.offset + mode) = 0.002 * std::cos(0.41 * (mode + 1));
        }
        for (const std::string_view name : {"gravity_gradient", "gravity_potential"}) {
            const Block &block = findBlock(modal.GetManifest().valueBlocks(), name);
            for (int i = 0; i < block.size; ++i)
                state(block.offset + i) = 0.01 * std::cos(0.13 * (i + 1));
        }

        const mfem::Vector nodalState     = liftState(modal, nodal, basis, state);
        const mfem::Vector nodalDirection = liftState(modal, nodal, basis, direction);
        prepare(modal, state, 1);
        prepare(nodal, nodalState, 1);
        checkClose(
            modal.GetPhysicalOperator().GetGeneratedVolumeDisplacement(),
            nodal.GetPhysicalOperator().GetGeneratedVolumeDisplacement(), 2.0e-12, 2.0e-14
        );

        const auto &modalDomain = modal.GetPhysicalOperator().GetDomainDeformation();
        const auto &nodalDomain = nodal.GetPhysicalOperator().GetDomainDeformation();
        CHECK_FALSE(modalDomain.hasNodalRadialParameters());
        CHECK(nodalDomain.hasNodalRadialParameters());
        const Block &nodalShape            = findBlock(nodal.GetManifest().valueBlocks(), "surface_deformation");
        const mfem::Vector modalParameters = readBlock(state, shape);
        const mfem::Vector nodalParameters = readBlock(nodalState, nodalShape);
        mfem::Vector modalTrace(modalDomain.surfaceDisplacementSize());
        mfem::Vector nodalTrace(nodalDomain.surfaceDisplacementSize());
        modalDomain.applySurfaceJacobian(modalParameters, readBlock(direction, shape), modalTrace);
        nodalDomain.applySurfaceJacobian(nodalParameters, readBlock(nodalDirection, nodalShape), nodalTrace);
        checkClose(modalTrace, nodalTrace, 3.0e-13, 2.0e-14);
        mfem::Vector traceDual(modalTrace.Size());
        for (int i = 0; i < traceDual.Size(); ++i)
            traceDual(i) = std::sin(0.23 * (i + 1));
        mfem::Vector modalDual(shape.size), nodalDual(nodalShape.size), expectedDual(shape.size);
        modalDomain.applySurfaceJacobianTranspose(modalParameters, traceDual, modalDual);
        nodalDomain.applySurfaceJacobianTranspose(nodalParameters, traceDual, nodalDual);
        basis.MultTranspose(nodalDual, expectedDual);
        checkClose(modalDual, expectedDual, 3.0e-13, 2.0e-13);

        mfem::Vector residual, nodalResidual, action, nodalAction;
        modal.BuildResidual(residual);
        nodal.BuildResidual(nodalResidual);
        modal.ApplyLinearization(direction, action);
        nodal.ApplyLinearization(nodalDirection, nodalAction);
        checkBlocks(residual, pullResidual(modal, nodal, basis, nodalResidual), modal.GetManifest().residualBlocks());
        checkBlocks(action, pullResidual(modal, nodal, basis, nodalAction), modal.GetManifest().residualBlocks());

        // A surface-only perturbation must also reach the retained field rows;
        // a full coupled comparison alone could overlook an inert shape block.
        mfem::Vector shapeDirection(direction.Size());
        shapeDirection = 0.0;
        writeBlock(shapeDirection, shape, readBlock(direction, shape));
        mfem::Vector shapeAction;
        modal.ApplyLinearization(shapeDirection, shapeAction);
        CHECK(
            readBlock(shapeAction, findBlock(modal.GetManifest().residualBlocks(), "poisson_balance")).Norml2() >
            1.0e-10
        );
        CHECK(readBlock(shapeAction, shapeRows).Norml2() > 1.0e-10);

        constexpr double step = 2.0e-5;
        mfem::Vector plus(state), minus(state), plusResidual, minusResidual;
        plus.Add(step, direction);
        minus.Add(-step, direction);
        prepare(modal, plus, 2);
        modal.BuildResidual(plusResidual);
        prepare(modal, minus, 3);
        modal.BuildResidual(minusResidual);
        plusResidual -= minusResidual;
        plusResidual /= 2.0 * step;
        checkBlocks(plusResidual, action, modal.GetManifest().residualBlocks(), 2.0e-5, 2.0e-9);

        auto normalized = normalization::makeNormalizedStellarEquilibriumOperator(modal);
        const mfem::Vector fixedStateFactors(normalized.GetNormalization().StateFactors());
        const mfem::Vector fixedResidualFactors(normalized.GetNormalization().ResidualFactors());
        REQUIRE(fixedStateFactors.Size() == modal.StateSize());
        REQUIRE(fixedResidualFactors.Size() == modal.EquationSize());
        for (int i = 0; i < fixedStateFactors.Size(); ++i) {
            REQUIRE(std::isfinite(fixedStateFactors(i)));
            REQUIRE(fixedStateFactors(i) > 0.0);
        }
        for (int i = 0; i < fixedResidualFactors.Size(); ++i) {
            REQUIRE(std::isfinite(fixedResidualFactors(i)));
            REQUIRE(fixedResidualFactors(i) > 0.0);
        }
        const auto prepareNormalized = [&](const mfem::Vector &physical, const std::uint64_t revision) {
            mfem::Vector coordinates;
            normalized.NormalizeState(physical, coordinates);
            if constexpr (GeneratedRotation) {
                (void)normalized.Prepare(coordinates, dependencies(revision));
            } else {
                (void)normalized.Prepare(coordinates, dependencies(revision), zeroRotation());
            }
            REQUIRE(normalized.IsPrepared());
        };
        prepareNormalized(state, 4);
        checkClose(normalized.GetPhysicalState(), state, 3.0e-14, 2.0e-14);
        mfem::Vector normalizedDirection, normalizedAction, expectedAction;
        normalized.NormalizeState(direction, normalizedDirection);
        normalized.Mult(normalizedDirection, normalizedAction);
        normalized.NormalizeResidual(action, expectedAction);
        checkBlocks(normalizedAction, expectedAction, modal.GetManifest().residualBlocks(), 3.0e-10, 2.0e-12);
        mfem::Vector normalizedResidual, expectedResidual;
        normalized.BuildResidual(normalizedResidual);
        normalized.NormalizeResidual(residual, expectedResidual);
        checkBlocks(normalizedResidual, expectedResidual, modal.GetManifest().residualBlocks(), 3.0e-10, 2.0e-12);

        // Repeated accepted-state-sized changes must not silently rotate or
        // rescale persistent harmonic coordinates through a current-shape metric.
        prepareNormalized(plus, 5);
        prepareNormalized(state, 6);
        for (int i = 0; i < fixedStateFactors.Size(); ++i)
            CHECK(normalized.GetNormalization().StateFactors()(i) == fixedStateFactors(i));
        for (int i = 0; i < fixedResidualFactors.Size(); ++i)
            CHECK(normalized.GetNormalization().ResidualFactors()(i) == fixedResidualFactors(i));
        mfem::Vector restoredResidual, restoredAction;
        modal.BuildResidual(restoredResidual);
        modal.ApplyLinearization(direction, restoredAction);
        checkBlocks(restoredResidual, residual, modal.GetManifest().residualBlocks());
        checkBlocks(restoredAction, action, modal.GetManifest().residualBlocks());
    }
} // namespace spherical_harmonic_equilibrium_test

TEST_CASE(
    "Native Harmonic Stellar Equilibrium Equals The Nodal Pullback Without A Central Constraint",
    tags::stellar_equilibrium_problem_integration & tags::spherical_harmonic_surface
) {
    int ranks = 0;
    MPI_Comm_size(MPI_COMM_WORLD, &ranks);
    if (ranks != 1) {
        SUCCEED("The current radial domain extension and harmonic prescription require one MPI rank.");
        return;
    }
    SECTION("Prescribed zero rotation") {
        spherical_harmonic_equilibrium_test::checkCoupledEquivalence<false>();
    }
    SECTION("Generated rotation from nonzero target angular momentum") {
        spherical_harmonic_equilibrium_test::checkCoupledEquivalence<true>();
    }
}

TEST_CASE(
    "Harmonic Degree And Azimuthal Cutoffs Determine Native Coupled Block Sizes",
    tags::stellar_equilibrium_problem_integration & tags::spherical_harmonic_surface
) {
    using namespace mean_field;
    using namespace spherical_harmonic_equilibrium_test;
    int ranks = 0;
    MPI_Comm_size(MPI_COMM_WORLD, &ranks);
    if (ranks != 1) {
        SUCCEED("The current radial domain extension and harmonic prescription require one MPI rank.");
        return;
    }
    struct Case final {
        int degree;
        std::optional<int> order;
        int count;
    };
    const std::array<Case, 3> cases{{{0, std::nullopt, 1}, {4, 0, 5}, {4, std::nullopt, 25}}};
    for (const Case &test : cases) {
        INFO("maximum degree = " << test.degree << ", expected parameters = " << test.count);
        const utils::Args args  = test_utils::setup_args();
        fem::FEM finiteElements = fem::setup_fem(args.mesh_file, args, 0);
        REQUIRE(finiteElements.okay());
        auto problem = equilibrium::discretize(
            makeModel<false>(),
            equilibrium::makeStellarDiscretization(
                std::move(finiteElements), normalization::PhysicalRieszDiagonal{dimensions::LengthValue{1.0}},
                deformation::SphericalHarmonicSurface{
                    {.maximumDegree = test.degree, .maximumAzimuthalOrder = test.order}
                }
            )
        );
        CHECK(findBlock(problem.GetManifest().valueBlocks(), "surface_deformation").size == test.count);
        CHECK(findBlock(problem.GetManifest().residualBlocks(), "surface_shape_balance").size == test.count);
        CHECK(problem.StateSize() == problem.EquationSize());
        auto normalized = normalization::makeNormalizedStellarEquilibriumOperator(problem);
        REQUIRE(normalized.GetNormalization().StateSize() == problem.StateSize());
        REQUIRE(normalized.GetNormalization().ResidualSize() == problem.EquationSize());
        for (const mfem::Vector *factors :
             {&normalized.GetNormalization().StateFactors(), &normalized.GetNormalization().ResidualFactors()}) {
            for (int i = 0; i < factors->Size(); ++i) {
                REQUIRE(std::isfinite((*factors)(i)));
                REQUIRE((*factors)(i) > 0.0);
            }
        }
    }
}