module;

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

#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::barotropic_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 displacementValue =
        mean_field::utils::blocks::get_value_block<Form>(mean_field::utils::blocks::displacement_field.geometry_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 displacementResidual = mean_field::utils::blocks::get_residual_block<Form>(
        mean_field::utils::blocks::displacement_field.geometry_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",
                                                                    "displacement",     "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), "Null-space 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 << "[null-space experiment] " << message << std::endl;
        }
    }

    [[nodiscard]] inline mean_field::operators::StellarEquilibriumDependencies make_dependencies() {
        return {
            .discretization    = {.identity = 2003, .revision = 1},
            .density           = {.identity = 2011, .revision = 1},
            .displacement      = {.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.displacement.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]] 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 unpinned_residual() const {
            const auto &layout                  = m_operator.GetLayout();
            const mfem::Vector reducedDensity   = const_value_view(m_currentState, layout, densityValue);
            const mfem::Vector displacement     = const_value_view(m_currentState, layout, displacementValue);
            const mfem::Vector gravityGradient  = const_value_view(m_currentState, layout, gravityGradientValue);
            const mfem::Vector gravityPotential = const_value_view(m_currentState, layout, gravityPotentialValue);
            const mfem::Vector gravityState =
                pack_gravity_state(reducedDensity, displacement, gravityGradient, gravityPotential);

            mfem::Vector gravity;
            mfem::Vector closure;
            mfem::Vector displacementRows;
            mfem::Vector hydrostatic;
            mfem::Vector mass;
            m_operator.GetGravityOperator().Mult(gravityState, gravity);
            m_operator.GetBarotropicClosureOperator().BuildResidual(closure);
            m_operator.GetDisplacementOperator().BuildResidual(displacementRows);
            m_operator.GetHydrostaticOperator().BuildResidual(hydrostatic);
            m_operator.GetSurfaceConstraintOperator().ApplyResidualRows(hydrostatic);
            m_operator.GetMassNormalizationOperator().BuildResidual(mass);

            return pack_residual(gravity, closure, displacementRows, hydrostatic, mass);
        }

        [[nodiscard]] mfem::Vector unpinned_jacobian_action(const mfem::Vector &direction) const {
            const auto &layout                           = m_operator.GetLayout();
            const mfem::Vector reducedDensityDirection   = const_value_view(direction, layout, densityValue);
            const mfem::Vector displacementDirection     = const_value_view(direction, layout, displacementValue);
            const mfem::Vector gravityGradientDirection  = const_value_view(direction, layout, gravityGradientValue);
            const mfem::Vector gravityPotentialDirection = const_value_view(direction, layout, gravityPotentialValue);
            const mfem::Vector reducedEnthalpyDirection  = const_value_view(direction, layout, enthalpyValue);
            const mfem::Vector bernoulliDirection        = const_value_view(direction, layout, bernoulliValue);

            const mfem::Vector gravityDirection          = pack_gravity_state(
                reducedDensityDirection, displacementDirection, gravityGradientDirection, gravityPotentialDirection
            );

            mfem::Vector gravity;
            mfem::Vector closure;
            mfem::Vector displacementRows;
            mfem::Vector hydrostatic;
            mfem::Vector mass;
            m_operator.GetGravityJacobianOperator().Mult(gravityDirection, gravity);
            m_operator.GetBarotropicClosureOperator().Mult(
                reducedDensityDirection, reducedEnthalpyDirection, displacementDirection, closure
            );
            m_operator.GetDisplacementOperator().ApplyCompleteJacobianAction(
                reducedDensityDirection, displacementDirection, gravityGradientDirection, reducedEnthalpyDirection,
                displacementRows
            );
            m_operator.GetHydrostaticOperator().ApplyCompleteJacobianAction(
                reducedEnthalpyDirection, gravityPotentialDirection, bernoulliDirection(0), displacementDirection,
                hydrostatic
            );
            m_operator.GetSurfaceConstraintOperator().ApplyJacobianRows(reducedEnthalpyDirection, hydrostatic);
            m_operator.GetMassNormalizationOperator().ApplyCompleteJacobianAction(
                reducedDensityDirection, displacementDirection, mass
            );

            return pack_residual(gravity, closure, displacementRows, hydrostatic, mass);
        }

    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 displacementTrue;
            mfem::Vector gravityGradientTrue;
            mfem::Vector gravityPotentialTrue;
            densityField.GetTrueDofs(densityTrue);
            enthalpyField.GetTrueDofs(enthalpyTrue);
            displacementField.GetTrueDofs(displacementTrue);
            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);

            assign_value_block(m_state, layout, densityValue, densityMap.gather(densityTrue));
            assign_value_block(m_state, layout, displacementValue, displacementTrue);
            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");
        }

        [[nodiscard]] mfem::Vector pack_residual(
            const mfem::Vector &gravity,
            const mfem::Vector &closure,
            const mfem::Vector &displacementRows,
            const mfem::Vector &hydrostatic,
            const mfem::Vector &mass
        ) const {
            const auto &layout = m_operator.GetLayout();
            mfem::Vector result(layout.residual_offsets().Last());
            result = 0.0;
            const mfem::Vector gravityGradient(gravity.GetData(), layout.size(gravityGradientResidual));
            const mfem::Vector gravityPotential(
                gravity.GetData() + layout.size(gravityGradientResidual), layout.size(gravityPotentialResidual)
            );
            residual_view(result, layout, gravityGradientResidual)  = gravityGradient;
            residual_view(result, layout, gravityPotentialResidual) = gravityPotential;
            residual_view(result, layout, densityResidual)          = closure;
            residual_view(result, layout, displacementResidual)     = displacementRows;
            residual_view(result, layout, enthalpyResidual)         = hydrostatic;
            residual_view(result, layout, massResidual)             = mass;
            return result;
        }

        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 RigidModeKind : std::uint8_t { translation, rotation };

    struct RigidMode final {
        std::string name;
        RigidModeKind kind;
        int axis;
        mfem::Vector direction;
    };

    [[nodiscard]] inline std::array<
        RigidMode,
        6>
    make_rigid_modes(const N3Equilibrium &fixture) {
        const auto &fem    = fixture.fem();
        const auto &layout = fixture.stellar_operator().GetLayout();
        std::array<RigidMode, 6> modes;

        for (int axis = 0; axis < 3; ++axis) {
            mfem::ParGridFunction translation(fem.displacementFes.get());
            mfem::Vector translationValue(3);
            translationValue       = 0.0;
            translationValue(axis) = 1.0;
            mfem::VectorConstantCoefficient coefficient(translationValue);
            translation.ProjectCoefficient(coefficient);
            mfem::Vector translationTrue;
            translation.GetTrueDofs(translationTrue);

            mfem::Vector direction(layout.value_offsets().Last());
            direction = 0.0;
            assign_value_block(direction, layout, displacementValue, translationTrue);
            modes[axis] = RigidMode{
                .name      = std::string("translation_") + static_cast<char>('x' + axis),
                .kind      = RigidModeKind::translation,
                .axis      = axis,
                .direction = std::move(direction)
            };
        }

        for (int axis = 0; axis < 3; ++axis) {
            mfem::ParGridFunction rotation(fem.displacementFes.get());
            mfem::VectorFunctionCoefficient coefficient(3, [axis](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(3);
                value            = 0.0;
                const int first  = (axis + 1) % 3;
                const int second = (axis + 2) % 3;
                value(first)     = -position(second);
                value(second)    = position(first);
            });
            rotation.ProjectCoefficient(coefficient);
            mfem::Vector rotationTrue;
            rotation.GetTrueDofs(rotationTrue);

            mfem::Vector direction(layout.value_offsets().Last());
            direction = 0.0;
            assign_value_block(direction, layout, displacementValue, rotationTrue);
            modes[3 + axis] = RigidMode{
                .name      = std::string("rotation_") + static_cast<char>('x' + axis),
                .kind      = RigidModeKind::rotation,
                .axis      = axis,
                .direction = std::move(direction)
            };
        }
        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, displacementResidual), communicator),
            global_norm(const_residual_view(action, layout, enthalpyResidual), communicator),
            global_norm(const_residual_view(action, layout, massResidual), communicator)
        };
    }
} // namespace experiment::null_space