module;

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <iostream>
#include <limits>
#include <string>
#include <utility>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>

export module experiment.stellar_null_space;

import mean_field;
import test_helpers;

export namespace experiment::null_space {
    using Form         = mean_field::utils::blocks::surface_deformed_stellar_equilibrium_form;
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using Model        = mean_field::models::StellarModel<mean_field::models::structure::PolytropicStructure>;

    constexpr auto densityValue =
        mean_field::utils::blocks::get_value_block<Form>(mean_field::utils::blocks::density_field.mass_term);
    constexpr auto surfaceDeformationValue = mean_field::utils::blocks::get_value_block<Form>(
        mean_field::utils::blocks::surface_deformation_field.parameters_term
    );
    constexpr auto gravityGradientValue =
        mean_field::utils::blocks::get_value_block<Form>(mean_field::utils::blocks::gravity_field.gradient_term);
    constexpr auto gravityPotentialValue =
        mean_field::utils::blocks::get_value_block<Form>(mean_field::utils::blocks::gravity_field.poisson_term);
    constexpr auto enthalpyValue =
        mean_field::utils::blocks::get_value_block<Form>(mean_field::utils::blocks::enthalpy_field.specific_term);
    constexpr auto bernoulliValue = mean_field::utils::blocks::get_value_block<Form>(
        mean_field::utils::blocks::barotropic_constant_field.mass_normalization_term
    );

    constexpr auto gravityGradientResidual =
        mean_field::utils::blocks::get_residual_block<Form>(mean_field::utils::blocks::gravity_field.gradient_term);
    constexpr auto gravityPotentialResidual =
        mean_field::utils::blocks::get_residual_block<Form>(mean_field::utils::blocks::gravity_field.poisson_term);
    constexpr auto densityResidual =
        mean_field::utils::blocks::get_residual_block<Form>(mean_field::utils::blocks::density_field.mass_term);
    constexpr auto surfaceShapeResidual = mean_field::utils::blocks::get_residual_block<Form>(
        mean_field::utils::blocks::surface_deformation_field.shape_equilibrium_term
    );
    constexpr auto enthalpyResidual =
        mean_field::utils::blocks::get_residual_block<Form>(mean_field::utils::blocks::enthalpy_field.specific_term);
    constexpr auto massResidual = mean_field::utils::blocks::get_residual_block<Form>(
        mean_field::utils::blocks::barotropic_constant_field.mass_normalization_term
    );

    inline constexpr std::array<const char *, 6> residualBlockNames{"gravity_gradient", "gravity_potential", "closure",
                                                                    "surface_shape",    "hydrostatic",       "mass"};

    template <int index>
    [[nodiscard]] mfem::Vector value_view(
        mfem::Vector &vector,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::value_block<index> block
    ) {
        return mfem::Vector(vector.GetData() + layout.offset(block), layout.size(block));
    }

    template <int index>
    [[nodiscard]] mfem::Vector const_value_view(
        const mfem::Vector &vector,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::value_block<index> block
    ) {
        return mfem::Vector(const_cast<mfem::real_t *>(vector.GetData()) + layout.offset(block), layout.size(block));
    }

    template <int index>
    [[nodiscard]] mfem::Vector residual_view(
        mfem::Vector &vector,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::residual_block<index> block
    ) {
        return mfem::Vector(vector.GetData() + layout.offset(block), layout.size(block));
    }

    template <int index>
    [[nodiscard]] mfem::Vector const_residual_view(
        const mfem::Vector &vector,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::residual_block<index> block
    ) {
        return mfem::Vector(const_cast<mfem::real_t *>(vector.GetData()) + layout.offset(block), layout.size(block));
    }

    template <int index>
    void assign_value_block(
        mfem::Vector &vector,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mean_field::utils::blocks::value_block<index> block,
        const mfem::Vector &source
    ) {
        MFEM_VERIFY(
            source.Size() == layout.size(block), "Surface-mode experiment received a block with the wrong size."
        );
        value_view(vector, layout, block) = source;
    }

    [[nodiscard]] inline double global_norm(
        const mfem::Vector &vector,
        const MPI_Comm communicator
    ) {
        const double localNormSquared = vector * vector;
        double globalNormSquared      = 0.0;
        MPI_Allreduce(&localNormSquared, &globalNormSquared, 1, MPI_DOUBLE, MPI_SUM, communicator);
        return std::sqrt(globalNormSquared);
    }

    inline void report_progress(
        const MPI_Comm communicator,
        const std::string &message
    ) {
        int rank = 0;
        MPI_Comm_rank(communicator, &rank);
        if (rank == 0) {
            std::cout << "[reduced-surface experiment] " << message << std::endl;
        }
    }

    [[nodiscard]] inline mean_field::operators::StellarEquilibriumDependencies make_dependencies() {
        return {
            .discretization     = {.identity = 2003, .revision = 1},
            .density            = {.identity = 2011, .revision = 1},
            .surfaceDeformation = {.identity = 2017, .revision = 1},
            .gravityGradient    = {.identity = 2027, .revision = 1},
            .gravityPotential   = {.identity = 2029, .revision = 1},
            .enthalpy           = {.identity = 2039, .revision = 1},
            .bernoulliConstant  = {.identity = 2053, .revision = 1},
            .rotation           = {.identity = 2063, .revision = 1},
            .targetMass         = {.identity = 2069, .revision = 1}
        };
    }

    inline void increment_state_revisions(mean_field::operators::StellarEquilibriumDependencies &dependencies) {
        ++dependencies.density.revision;
        ++dependencies.surfaceDeformation.revision;
        ++dependencies.gravityGradient.revision;
        ++dependencies.gravityPotential.revision;
        ++dependencies.enthalpy.revision;
        ++dependencies.bernoulliConstant.revision;
    }

    [[nodiscard]] inline mfem::Vector pack_gravity_state(
        const mfem::Vector &density,
        const mfem::Vector &displacement,
        const mfem::Vector &gravityGradient,
        const mfem::Vector &gravityPotential
    ) {
        const std::array<int, 5> offsets{
            0, density.Size(), density.Size() + displacement.Size(),
            density.Size() + displacement.Size() + gravityGradient.Size(),
            density.Size() + displacement.Size() + gravityGradient.Size() + gravityPotential.Size()
        };
        mfem::Vector packed(offsets.back());
        mfem::Vector(packed.GetData() + offsets[0], density.Size())          = density;
        mfem::Vector(packed.GetData() + offsets[1], displacement.Size())     = displacement;
        mfem::Vector(packed.GetData() + offsets[2], gravityGradient.Size())  = gravityGradient;
        mfem::Vector(packed.GetData() + offsets[3], gravityPotential.Size()) = gravityPotential;
        return packed;
    }

    [[nodiscard]] inline Model make_model() {
        const double pi                    = std::acos(-1.0);
        const double targetMass            = mean_field::utils::MASS;
        constexpr double dimensionlessMass = 2.0182359509662283534;
        const double polytropicConstant =
            pi * mean_field::utils::G * std::pow(targetMass / (4.0 * pi * dimensionlessMass), 2.0 / 3.0);

        return Model{
            mean_field::models::structure::PolytropicStructure{
                mean_field::eos::Polytrope{3.0, polytropicConstant}, targetMass
            },
            mean_field::surface::ConstantPressureSurface{mean_field::eos::PressureValue{0.0}}
        };
    }

    class N3Equilibrium final {
    public:
        explicit N3Equilibrium(mean_field::utils::Args args)
            : m_args(std::move(args)),
              m_fem(
                  mean_field::fem::setup_fem(
                      m_args.mesh_file,
                      m_args,
                      0
                  )
              ),
              m_model(make_model()),
              m_operator(
                  m_fem,
                  *m_fem.domainMapperStateless,
                  m_model
              ),
              m_state(m_operator.GetLayout().value_offsets().Last()),
              m_dependencies(make_dependencies()) {
            MFEM_VERIFY(m_fem.okay(), "The null-space experiment could not construct the finite-element problem.");
            m_state = 0.0;
            initialize_state();
        }

        [[nodiscard]] mean_field::fem::FEM &fem() noexcept {
            return m_fem;
        }

        [[nodiscard]] const mean_field::fem::FEM &fem() const noexcept {
            return m_fem;
        }

        [[nodiscard]] Model &model() noexcept {
            return m_model;
        }

        [[nodiscard]] mean_field::operators::PreparedStellarEquilibriumOperator &stellar_operator() noexcept {
            return m_operator;
        }

        [[nodiscard]] const mean_field::operators::PreparedStellarEquilibriumOperator &
        stellar_operator() const noexcept {
            return m_operator;
        }

        [[nodiscard]] const mfem::Vector &state() const noexcept {
            return m_state;
        }

        [[nodiscard]] mean_field::physics::RigidRotation rotation(const double fractionOfKeplerian) const {
            const double radius         = mean_field::utils::RADIUS;
            const double mass           = mean_field::utils::MASS;
            const double keplerianSpeed = std::sqrt(mean_field::utils::G * mass / (radius * radius * radius));

            mfem::Vector angularVelocity(3);
            angularVelocity    = 0.0;
            angularVelocity(2) = fractionOfKeplerian * keplerianSpeed;

            mfem::Vector center(3);
            center = 0.0;
            return mean_field::physics::RigidRotation(angularVelocity, center);
        }

        void prepare(
            const mfem::Vector &state,
            const mean_field::physics::RigidRotation &rotation
        ) {
            m_currentState = state;
            increment_state_revisions(m_dependencies);
            ++m_dependencies.rotation.revision;
            m_operator.Prepare(state, m_dependencies, rotation);
        }

        [[nodiscard]] mfem::Vector residual() const {
            mfem::Vector result;
            m_operator.BuildResidual(result);
            return result;
        }

        [[nodiscard]] mfem::Vector jacobian_action(const mfem::Vector &direction) const {
            mfem::Vector result;
            m_operator.Mult(direction, result);
            return result;
        }

        [[nodiscard]] mfem::Vector lifted_surface_direction(const mfem::Vector &rootDirection) const {
            const auto &layout                  = m_operator.GetLayout();
            const mfem::Vector surfaceDirection = const_value_view(rootDirection, layout, surfaceDeformationValue);
            mfem::Vector volumeDirection(m_operator.GetDomainDeformation().volumeDisplacementSize());
            m_operator.GetDomainDeformation().applyJacobian(
                m_operator.GetSurfaceDeformationParameters(), surfaceDirection, volumeDirection
            );
            return volumeDirection;
        }

    private:
        void initialize_state() {
            report_progress(m_fem.mesh->GetComm(), "constructing the analytic n=3 Lane-Emden state");

            constexpr double surfaceCoordinate = 6.8968486193769603755;
            constexpr int radialSampleCount    = 8192;
            const double pi                    = std::acos(-1.0);
            const double radius                = mean_field::utils::RADIUS;
            const double targetMass            = mean_field::utils::MASS;
            constexpr double dimensionlessMass = 2.0182359509662283534;
            const double polytropicConstant =
                pi * mean_field::utils::G * std::pow(targetMass / (4.0 * pi * dimensionlessMass), 2.0 / 3.0);
            const double centralDensity =
                std::pow(surfaceCoordinate * std::sqrt(polytropicConstant / (pi * mean_field::utils::G)) / radius, 3.0);

            const mean_field::models::structure::StructureSeed seed =
                m_model.makeInitialSeed({.centralDensity = centralDensity, .radialSampleCount = radialSampleCount});

            const auto interpolate = [](const mfem::Vector &radii, const mfem::Vector &values, const double r) {
                if (r <= radii(0)) {
                    return values(0);
                }
                const int finalIndex = radii.Size() - 1;
                if (r >= radii(finalIndex)) {
                    return values(finalIndex);
                }
                int lower = 0;
                int upper = finalIndex;
                while (upper - lower > 1) {
                    const int middle = lower + (upper - lower) / 2;
                    if (radii(middle) <= r) {
                        lower = middle;
                    } else {
                        upper = middle;
                    }
                }
                const double fraction = (r - radii(lower)) / (radii(upper) - radii(lower));
                return (1.0 - fraction) * values(lower) + fraction * values(upper);
            };

            mfem::FunctionCoefficient densityCoefficient([&seed, &interpolate](const mfem::Vector &position) {
                const double r = position.Norml2();
                return r >= seed.stellarRadius ? 0.0 : interpolate(seed.radius, seed.density, r);
            });
            mfem::FunctionCoefficient enthalpyCoefficient([&seed, &interpolate](const mfem::Vector &position) {
                const double r = position.Norml2();
                return r >= seed.stellarRadius ? 0.0 : interpolate(seed.radius, seed.enthalpy, r);
            });

            mfem::ParGridFunction densityField(m_fem.densityFes.get());
            mfem::ParGridFunction enthalpyField(m_fem.enthalpyFes.get());
            mfem::ParGridFunction displacementField(m_fem.displacementFes.get());
            densityField      = 0.0;
            enthalpyField     = 0.0;
            displacementField = 0.0;
            densityField.ProjectCoefficient(densityCoefficient);
            enthalpyField.ProjectCoefficient(enthalpyCoefficient);
            *m_fem.displacement = displacementField;

            report_progress(m_fem.mesh->GetComm(), "solving the gravity field for the seed state");
            const mean_field::physics::GravitySolution gravity =
                mean_field::physics::solve_gravity_field(m_fem, m_args, densityField, displacementField);

            mfem::Vector densityTrue;
            mfem::Vector enthalpyTrue;
            mfem::Vector gravityGradientTrue;
            mfem::Vector gravityPotentialTrue;
            densityField.GetTrueDofs(densityTrue);
            enthalpyField.GetTrueDofs(enthalpyTrue);
            gravity.gradPhi.GetTrueDofs(gravityGradientTrue);
            gravity.phi.GetTrueDofs(gravityPotentialTrue);

            const auto &layout = m_operator.GetLayout();
            const mean_field::field::FieldDofMap densityMap =
                mean_field::field::make_field_dof_map<mean_field::field::Density, DomainSchema>(*m_fem.densityFes);
            const mean_field::field::FieldDofMap enthalpyMap =
                mean_field::field::make_field_dof_map<mean_field::field::Enthalpy, DomainSchema>(*m_fem.enthalpyFes);

            mfem::Vector surfaceParameters(layout.size(surfaceDeformationValue));
            surfaceParameters = 0.0;
            assign_value_block(m_state, layout, densityValue, densityMap.gather(densityTrue));
            assign_value_block(m_state, layout, surfaceDeformationValue, surfaceParameters);
            assign_value_block(m_state, layout, gravityGradientValue, gravityGradientTrue);
            assign_value_block(m_state, layout, gravityPotentialValue, gravityPotentialTrue);
            assign_value_block(m_state, layout, enthalpyValue, enthalpyMap.gather(enthalpyTrue));
            value_view(m_state, layout, bernoulliValue)(0) = -mean_field::utils::G * targetMass / radius;

            m_currentState                                 = m_state;
            prepare(m_state, rotation(0.0));
            report_progress(m_fem.mesh->GetComm(), "analytic state is prepared");
        }

        mean_field::utils::Args m_args;
        mean_field::fem::FEM m_fem;
        Model m_model;
        mean_field::operators::PreparedStellarEquilibriumOperator m_operator;
        mfem::Vector m_state;
        mfem::Vector m_currentState;
        mean_field::operators::StellarEquilibriumDependencies m_dependencies;
    };

    enum class SurfaceModeKind : std::uint8_t {
        uniform_radial,
        translation_like_dipole,
        oblate_quadrupole,
        spherical_harmonic
    };

    struct SurfaceMode final {
        std::string name;
        SurfaceModeKind kind;
        int axis;
        mfem::Vector direction;
    };

    [[nodiscard]] inline const char *surface_mode_kind_name(const SurfaceModeKind kind) noexcept {
        switch (kind) {
        case SurfaceModeKind::uniform_radial:
            return "uniform_radial";
        case SurfaceModeKind::translation_like_dipole:
            return "translation_like_dipole";
        case SurfaceModeKind::oblate_quadrupole:
            return "oblate_quadrupole";
        case SurfaceModeKind::spherical_harmonic:
            return "spherical_harmonic";
        }
        return "unknown";
    }

    [[nodiscard]] inline double zonal_legendre(
        const int degree,
        const double cosineOfPolarAngle
    ) {
        MFEM_VERIFY(degree >= 0, "A zonal spherical-harmonic degree must be non-negative.");
        const double coordinate = std::clamp(cosineOfPolarAngle, -1.0, 1.0);
        if (degree == 0) {
            return 1.0;
        }
        if (degree == 1) {
            return coordinate;
        }

        double previousPrevious = 1.0;
        double previous         = coordinate;
        for (int order = 2; order <= degree; ++order) {
            const double current = ((2.0 * static_cast<double>(order) - 1.0) * coordinate * previous -
                                    (static_cast<double>(order) - 1.0) * previousPrevious) /
                                   static_cast<double>(order);
            previousPrevious     = previous;
            previous             = current;
        }
        return previous;
    }

    [[nodiscard]] inline std::vector<SurfaceMode> make_surface_modes(N3Equilibrium &fixture) {
        const auto &layout  = fixture.stellar_operator().GetLayout();
        auto deformation    = fixture.model().compileDomainDeformation(fixture.fem());
        const auto &surface = deformation.surfaceDeformationPrescription();
        MFEM_VERIFY(
            surface.parameterCount() == layout.size(surfaceDeformationValue),
            "The diagnostic surface prescription does not match the root surface block."
        );

        const auto make_root_direction = [&layout](const mfem::Vector &surfaceDirection) {
            mfem::Vector direction(layout.value_offsets().Last());
            direction = 0.0;
            assign_value_block(direction, layout, surfaceDeformationValue, surfaceDirection);
            return direction;
        };

        std::vector<SurfaceMode> modes;
        modes.reserve(6);

        mfem::Vector uniform(surface.parameterCount());
        for (int parameter = 0; parameter < uniform.Size(); ++parameter) {
            uniform(parameter) = surface.referenceRadius(parameter);
        }
        modes.push_back(
            {.name      = "uniform_radial_homology",
             .kind      = SurfaceModeKind::uniform_radial,
             .axis      = -1,
             .direction = make_root_direction(uniform)}
        );

        for (int axis = 0; axis < surface.spatialDimension(); ++axis) {
            mfem::Vector dipole(surface.parameterCount());
            for (int parameter = 0; parameter < dipole.Size(); ++parameter) {
                dipole(parameter) = surface.radialDirection(parameter, axis);
            }
            modes.push_back(
                {.name      = std::string("translation_like_dipole_") + static_cast<char>('x' + axis),
                 .kind      = SurfaceModeKind::translation_like_dipole,
                 .axis      = axis,
                 .direction = make_root_direction(dipole)}
            );
        }

        mfem::Vector quadrupole(surface.parameterCount());
        for (int parameter = 0; parameter < quadrupole.Size(); ++parameter) {
            const double polarDirection = surface.radialDirection(parameter, 2);
            quadrupole(parameter) = surface.referenceRadius(parameter) * (1.0 - 3.0 * polarDirection * polarDirection);
        }
        modes.push_back(
            {.name      = "axisymmetric_oblate_quadrupole_z",
             .kind      = SurfaceModeKind::oblate_quadrupole,
             .axis      = 2,
             .direction = make_root_direction(quadrupole)}
        );

        constexpr int diagnosticAngularDegree = 12;
        mfem::Vector sphericalHarmonic(surface.parameterCount());
        double localMaximumMagnitude = 0.0;
        for (int parameter = 0; parameter < sphericalHarmonic.Size(); ++parameter) {
            const double angularValue = zonal_legendre(diagnosticAngularDegree, surface.radialDirection(parameter, 2));
            sphericalHarmonic(parameter) = surface.referenceRadius(parameter) * angularValue;
            localMaximumMagnitude        = std::max(localMaximumMagnitude, std::abs(angularValue));
        }
        double globalMaximumMagnitude = 0.0;
        MPI_Allreduce(
            &localMaximumMagnitude, &globalMaximumMagnitude, 1, MPI_DOUBLE, MPI_MAX, fixture.fem().mesh->GetComm()
        );
        MFEM_VERIFY(globalMaximumMagnitude > 0.0, "The spherical-harmonic surface mode has zero amplitude.");
        sphericalHarmonic /= globalMaximumMagnitude;
        modes.push_back(
            {.name      = "zonal_spherical_harmonic_l12",
             .kind      = SurfaceModeKind::spherical_harmonic,
             .axis      = -1,
             .direction = make_root_direction(sphericalHarmonic)}
        );

        return modes;
    }

    [[nodiscard]] inline std::array<
        double,
        6>
    residual_block_norms(
        const mfem::Vector &action,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const MPI_Comm communicator
    ) {
        return {
            global_norm(const_residual_view(action, layout, gravityGradientResidual), communicator),
            global_norm(const_residual_view(action, layout, gravityPotentialResidual), communicator),
            global_norm(const_residual_view(action, layout, densityResidual), communicator),
            global_norm(const_residual_view(action, layout, surfaceShapeResidual), communicator),
            global_norm(const_residual_view(action, layout, enthalpyResidual), communicator),
            global_norm(const_residual_view(action, layout, massResidual), communicator)
        };
    }
} // namespace experiment::null_space