module;
#include "mfem.hpp"
#include <array>
#include <cmath>
#include <format>
#include <source_location>
#include <string_view>
#include <unordered_map>

module mean_field;
import :mapping.coefficients;
import :analysis.integral;

namespace {
    double centrifugal_potential(
        const mfem::Vector &phys_x,
        const double omega
    ) {
        const double s2 = std::pow(phys_x(0), 2) + std::pow(phys_x(1), 2);
        return -0.5 * s2 * std::pow(omega, 2);
    }

    void grid_function_to_true_dofs(
        const mfem::ParFiniteElementSpace &finite_element_space,
        const mfem::GridFunction &grid_function,
        mfem::Vector &true_dofs
    ) {
        MFEM_VERIFY(
            grid_function.Size() == finite_element_space.GetVSize(),
            "The grid function does not match the requested finite-element "
            "space."
        );

        true_dofs.SetSize(finite_element_space.GetTrueVSize());

        const mfem::Operator *restriction = finite_element_space.GetRestrictionMatrix();

        if (restriction != nullptr) {
            restriction->Mult(grid_function, true_dofs);
        } else {
            MFEM_VERIFY(
                grid_function.Size() == true_dofs.Size(), "A finite-element space without a restriction operator must "
                                                          "have "
                                                          "matching local and true sizes."
            );

            true_dofs = grid_function;
        }
    }
} // namespace

namespace mean_field::physics {
    GravitySolution grav_potential(
        fem::FEM &f,
        const utils::Args &args,
        const mfem::GridFunction &rho,
        const bool phi_warm
    ) {
        MFEM_VERIFY(
            f.densityFes != nullptr && rho.FESpace() == f.densityFes.get(),
            "Gravity solve requires rho to use the registered density space."
        );
        MFEM_VERIFY(f.gravityPotentialFes != nullptr, "Gravity solve requires the registered gravity-potential space.");

        mfem::Array<int> outer_bdr_marker(f.mesh->bdr_attributes.Max());
        outer_bdr_marker    = 0;
        outer_bdr_marker[1] = 1;

        mfem::ParLinearForm g_rhs(f.gravityFluxFes.get());

        // ReSharper disable once CppTooWideScope
        std::unique_ptr<mfem::Coefficient> boundary_potential_coeff;

        if (!f.has_mapping()) { // We only need to explicitly add a boundary
                                // integrator if a mapping is not being used. In
                                // the case where the outer domain has been
                                // compactified the φ=0 boundary condition is
                                // the natural condition and MFEM automatically
                                // handles this
            auto boundary_potential = [&f](const mfem::Vector &x_physical) {
                return l2_multipole_potential(f, utils::MASS, x_physical);
            };

            boundary_potential_coeff = std::make_unique<mfem::FunctionCoefficient>(boundary_potential);
            auto boundary_integrator =
                std::make_unique<mfem::VectorFEBoundaryFluxLFIntegrator>(*boundary_potential_coeff);
            const mfem::FiniteElement &boundary_element = *f.gravityFluxFes->GetTypicalTraceElement();

            f.quadratureFactory->configure_gravity_boundary(
                *boundary_integrator, quadrature::QuadratureRole::discretization, boundary_element,
                utils::DOMAINS::VACUUM, quadrature::MappingKind::none
            );
            g_rhs.AddBoundaryIntegrator(boundary_integrator.release(), outer_bdr_marker);
        }

        g_rhs.Assemble();
        mfem::GridFunctionCoefficient rho_coeff(&rho);
        mfem::ConstantCoefficient G4pi(4.0 * M_PI * utils::G);
        mfem::ProductCoefficient source_coeff(G4pi, rho_coeff);
        mfem::ParLinearForm f_rhs(f.gravityPotentialFes.get());

        std::unique_ptr<mfem::Coefficient> mapped_source_coeff;
        mfem::Coefficient *active_source_coeff      = &source_coeff;
        quadrature::MappingKind source_mapping_kind = quadrature::MappingKind::none;

        if (f.has_mapping()) {
            mapped_source_coeff = std::make_unique<mapping::MappedScalarCoefficient>(*f.mapping, source_coeff);
            active_source_coeff = mapped_source_coeff.get();
            source_mapping_kind = quadrature::MappingKind::general;
        }

        auto source_integrator = std::make_unique<mfem::DomainLFIntegrator>(*active_source_coeff);
        const mfem::FiniteElement &source_test_element           = *f.gravityPotentialFes->GetTypicalFE();
        const mfem::ElementTransformation &source_transformation = *f.mesh->GetElementTransformation(0);
        const int source_coefficient_order                       = f.densityFes->GetMaxElementOrder();

        f.quadratureFactory->configure_gravity_source(
            *source_integrator, quadrature::QuadratureRole::discretization, source_test_element, source_transformation,
            source_coefficient_order, utils::DOMAINS::STELLAR, source_mapping_kind
        );
        f_rhs.AddDomainIntegrator(source_integrator.release(), f.gravityContext.stellar_mask);
        f_rhs.Assemble();

        mfem::BlockVector RHS(f.gravityBlockTrueOffsets);
        RHS.GetBlock(0) = *g_rhs.ParallelAssemble();
        RHS.GetBlock(1) = *f_rhs.ParallelAssemble();

        mfem::BlockVector X(f.gravityBlockTrueOffsets);
        X = 0.0;
        f.gravityContext.minres->SetOperator(*f.gravityContext.block_A);
        f.gravityContext.minres->Mult(RHS, X);

        GravitySolution solution(f);
        solution.gradPhi.SetFromTrueDofs(X.GetBlock(0));
        solution.phi.SetFromTrueDofs(X.GetBlock(1));

        return solution;
    }

    mfem::GridFunction get_potential(
        fem::FEM &fem,
        const utils::Args &args,
        const mfem::GridFunction &rho,
        const bool warm
    ) {
        auto phi = grav_potential(fem, args, rho, warm);

        if (args.r.enabled) {
            auto rot = [&fem, &args](const mfem::Vector &x) {
                mfem::Vector rel_x = x;
                rel_x -= fem.com;
                return centrifugal_potential(rel_x, args.r.omega);
            };

            std::unique_ptr<mfem::Coefficient> centrifugal_coeff;
            if (fem.has_mapping()) {
                centrifugal_coeff = std::make_unique<mapping::PhysicalPositionFunctionCoefficient>(*fem.mapping, rot);
            } else {
                centrifugal_coeff = std::make_unique<mfem::FunctionCoefficient>(rot);
            }

            mfem::GridFunction centrifugal_gf(fem.gravityPotentialFes.get());
            centrifugal_gf.ProjectCoefficient(*centrifugal_coeff);

            phi.phi += centrifugal_gf;
        }
        return phi.phi;
    }

    mfem::DenseMatrix compute_quadrupole_moment_tensor(
        const fem::FEM &fem,
        const mfem::GridFunction &rho,
        const mfem::Vector &com
    ) {
        const int dim = fem.mesh->Dimension();
        mfem::DenseMatrix local_Q(dim, dim);
        local_Q = 0.0;

        for (int i = 0; i < fem.mesh->GetNE(); ++i) {
            if (fem.mesh->GetAttribute(i) == 3)
                continue;

            mfem::ElementTransformation *trans = fem.mesh->GetElementTransformation(i);
            using DensityField                 = field::Field<field::Density>;
            const quadrature::Query query      = DensityField::make_query<field::Density::Form::Quadrupole>(
                quadrature::QuadratureRole::diagnostic, trans->OrderW(), std::array<int, 1>{2}, utils::DOMAINS::STELLAR,
                fem.has_mapping() ? quadrature::MappingKind::general : quadrature::MappingKind::none
            );
            const mfem::IntegrationRule &ir =
                *fem.quadratureFactory->get(query, trans->GetGeometryType()).integration_rule;

            for (int j = 0; j < ir.GetNPoints(); ++j) {
                const mfem::IntegrationPoint &ip = ir.IntPoint(j);
                trans->SetIntPoint(&ip);

                double weight = trans->Weight() * ip.weight;

                if (fem.has_mapping()) {
                    weight *= fem.mapping->ComputeDetJ(*trans, ip);
                }

                const double rho_val = rho.GetValue(i, ip);

                mfem::Vector phys_point(dim);
                if (fem.has_mapping()) {
                    fem.mapping->GetPhysicalPoint(*trans, ip, phys_point);
                } else {
                    trans->Transform(ip, phys_point);
                }

                mfem::Vector x_prime(dim);
                double r_sq = 0.0;

                for (int d = 0; d < dim; ++d) {
                    x_prime(d) = phys_point(d) - com(d);
                    r_sq += x_prime(d) * x_prime(d);
                }

                for (int m = 0; m < dim; ++m) {
                    for (int n = 0; n < dim; ++n) {
                        const double delta   = (m == n) ? 1.0 : 0.0;
                        const double contrib = 3.0 * x_prime(m) * x_prime(n) - delta * r_sq;
                        local_Q(m, n) += rho_val * contrib * weight;
                    }
                }
            }
        }

        mfem::DenseMatrix global_Q(dim, dim);
        MPI_Allreduce(local_Q.GetData(), global_Q.GetData(), dim * dim, MPI_DOUBLE, MPI_SUM, fem.mesh->GetComm());

        return global_Q;
    }

    double l2_multipole_potential(
        const fem::FEM &fem,
        const double total_mass,
        const mfem::Vector &phys_x
    ) {
        const double r = phys_x.Norml2();
        if (r < 1e-12)
            return 0.0;

        const int dim = fem.mesh->Dimension();

        mfem::Vector n(phys_x);
        n /= r;

        double l2_mult_factor = 0.0;
        for (int i = 0; i < dim; ++i) {
            for (int j = 0; j < dim; ++j) {
                l2_mult_factor += fem.Q(i, j) * n(i) * n(j);
            }
        }

        const double l2_contrib = -(utils::G / (2.0 * std::pow(r, 3))) * l2_mult_factor;

        const double l0_contrib = -utils::G * total_mass / r;

        // l1 contribution is zero for a system centered on its COM
        return l0_contrib + l2_contrib;
    }

    void update_stiffness_matrix(fem::FEM &f) {
        mfem::Array<int> empty_tdofs;

        // ==========================================
        // 1. Partially Assemble the High-Order Mass Block
        // ==========================================
        f.gravityContext.m_form = std::make_unique<mfem::ParBilinearForm>(f.gravityFluxFes.get());
        f.gravityContext.m_form->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);

        std::unique_ptr<mfem::VectorFEMassIntegrator> hdiv_mass_integrator;

        if (f.has_mapping()) {
            f.gravityContext.mapped_hdiv_mass_coeff =
                std::make_unique<mapping::MappedHDivMassCoefficient>(*f.mapping, f.mesh->Dimension());
            hdiv_mass_integrator =
                std::make_unique<mfem::VectorFEMassIntegrator>(*f.gravityContext.mapped_hdiv_mass_coeff);
        } else {
            f.gravityContext.mapped_hdiv_mass_coeff.reset();
            hdiv_mass_integrator = std::make_unique<mfem::VectorFEMassIntegrator>();
        }

        const mfem::FiniteElement &hdiv_element                = *f.gravityFluxFes->GetTypicalFE();
        const mfem::ElementTransformation &hdiv_transformation = *f.mesh->GetElementTransformation(0);
        const quadrature::MappingKind mapping_kind =
            f.has_mapping() ? quadrature::MappingKind::general : quadrature::MappingKind::none;

        f.quadratureFactory->configure_gravity_hdiv_mass(
            *hdiv_mass_integrator, quadrature::QuadratureRole::discretization, hdiv_element, hdiv_transformation,
            utils::DOMAINS::ALL, mapping_kind
        );
        f.gravityContext.m_form->AddDomainIntegrator(hdiv_mass_integrator.release());

        f.gravityContext.m_form->Assemble();

        // ==========================================
        // 2. Partially Assemble the High-Order Divergence Block
        // ==========================================
        f.gravityContext.b_form =
            std::make_unique<mfem::ParMixedBilinearForm>(f.gravityFluxFes.get(), f.gravityPotentialFes.get());
        f.gravityContext.b_form->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);

        auto divergence_discretization_integrator = std::make_unique<mfem::VectorFEDivergenceIntegrator>();
        const mfem::FiniteElement &divergence_discretization_test_element = *f.gravityPotentialFes->GetTypicalFE();

        f.quadratureFactory->configure_gravity_divergence(
            *divergence_discretization_integrator, quadrature::QuadratureRole::discretization, hdiv_element,
            divergence_discretization_test_element, hdiv_transformation, utils::DOMAINS::ALL,
            quadrature::MappingKind::none
        );
        f.gravityContext.b_form->AddDomainIntegrator(divergence_discretization_integrator.release());

        f.gravityContext.b_form->Assemble();

        MFEM_VERIFY(
            f.domainMapperStateless != nullptr, "Gravity source partial assembly requires the stateless domain "
                                                "mapper."
        );

        mfem::Vector displacement_true(f.displacementFes->GetTrueVSize());
        displacement_true                             = 0.0;

        const mfem::GridFunction *active_displacement = f.mapping->GetDisplacement();

        if (active_displacement != nullptr) {
            grid_function_to_true_dofs(*f.displacementFes, *active_displacement, displacement_true);
        }

        auto source_form =
            std::make_unique<operators::PreparedMappedGravitySourceOperator>(f, *f.domainMapperStateless);

        source_form->Prepare(displacement_true);

        f.gravityContext.source_form = std::move(source_form);
        // ==========================================
        // 3. Assemble Global Block Operator
        // ==========================================
        f.gravityContext.BT          = std::make_unique<mfem::TransposeOperator>(f.gravityContext.b_form.get());

        f.gravityContext.block_A     = std::make_unique<mfem::BlockOperator>(f.gravityBlockTrueOffsets);
        f.gravityContext.block_A->SetBlock(0, 0, f.gravityContext.m_form.get());
        f.gravityContext.block_A->SetBlock(0, 1, f.gravityContext.BT.get());
        f.gravityContext.block_A->SetBlock(1, 0, f.gravityContext.b_form.get());

        // ==========================================
        // 4. Construct a mapped Schur preconditioner
        // ==========================================
        mfem::Vector mass_diagonal(f.gravityFluxFes->GetTrueVSize());
        f.gravityContext.m_form->AssembleDiagonal(mass_diagonal);

        mfem::Vector inverse_mass_diagonal(mass_diagonal);

        for (int i = 0; i < inverse_mass_diagonal.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(inverse_mass_diagonal(i)) && inverse_mass_diagonal(i) > 0.0,
                "Mapped RT mass matrix has a non-positive or non-finite "
                "diagonal "
                "entry."
            );
            inverse_mass_diagonal(i) = 1.0 / inverse_mass_diagonal(i);
        }

        mfem::ParMixedBilinearForm b_preconditioner(f.gravityFluxFes.get(), f.gravityPotentialFes.get());
        auto divergence_preconditioner_integrator           = std::make_unique<mfem::VectorFEDivergenceIntegrator>();

        const mfem::FiniteElement &divergence_trial_element = *f.gravityFluxFes->GetTypicalFE();
        const mfem::FiniteElement &divergence_test_element  = *f.gravityPotentialFes->GetTypicalFE();
        const mfem::ElementTransformation &divergence_transformation = *f.mesh->GetElementTransformation(0);

        f.quadratureFactory->configure_gravity_divergence(
            *divergence_preconditioner_integrator, quadrature::QuadratureRole::preconditioner, divergence_trial_element,
            divergence_test_element, divergence_transformation, utils::DOMAINS::ALL, quadrature::MappingKind::none
        );
        b_preconditioner.AddDomainIntegrator(divergence_preconditioner_integrator.release());
        b_preconditioner.Assemble();
        b_preconditioner.Finalize();
        std::unique_ptr<mfem::HypreParMatrix> b_matrix(b_preconditioner.ParallelAssemble());
        std::unique_ptr<mfem::HypreParMatrix> inverse_mass_b_transpose(b_matrix->Transpose());

        inverse_mass_b_transpose->ScaleRows(inverse_mass_diagonal);

        f.gravityContext.Schur.reset(mfem::ParMult(b_matrix.get(), inverse_mass_b_transpose.get()));

        // ==========================================
        // 5. Wire Up the preconditioners
        // ==========================================
        f.gravityContext.prec_M = std::make_unique<mfem::OperatorJacobiSmoother>(mass_diagonal, empty_tdofs);
        f.gravityContext.prec_Phi->SetOperator(*f.gravityContext.Schur);
        f.gravityContext.block_prec->SetDiagonalBlock(0, f.gravityContext.prec_M.get());
        f.gravityContext.block_prec->SetDiagonalBlock(1, f.gravityContext.prec_Phi.get());
    }

    GravitySolution grav_potential_new(
        fem::FEM &f,
        const utils::Args &args,
        const mfem::GridFunction &rho,
        const mfem::GridFunction &displacement
    ) {
        MFEM_VERIFY(f.mesh != nullptr, "Gravity initialization requires a parallel mesh.");
        MFEM_VERIFY(f.densityFes != nullptr, "Gravity initialization requires the density finite-element space.");
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "Gravity initialization requires the gravity-potential "
                                              "finite-element "
                                              "space."
        );
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "Gravity initialization requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "Gravity initialization requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(f.domainMapperStateless != nullptr, "Gravity initialization requires the stateless domain mapper.");
        MFEM_VERIFY(f.gravityContext.b_form != nullptr, "Gravity initialization requires the divergence operator.");
        MFEM_VERIFY(
            f.gravityContext.BT != nullptr, "Gravity initialization requires the transpose divergence operator."
        );
        MFEM_VERIFY(
            f.gravityContext.block_prec != nullptr, "Gravity initialization requires the gravity block preconditioner."
        );
        MFEM_VERIFY(
            rho.FESpace() == f.densityFes.get(), "Gravity initialization requires density to use the FEM density "
                                                 "space."
        );
        MFEM_VERIFY(
            displacement.FESpace() == f.displacementFes.get(),
            "Gravity initialization requires displacement to use the FEM "
            "Vec_H1 "
            "space."
        );

        using form = utils::blocks::gravity_field_form;

        constexpr auto gravity_gradient_residual_block =
            utils::blocks::get_residual_block<form>(utils::blocks::gravity_field.gradient_term);

        constexpr auto gravity_poisson_residual_block =
            utils::blocks::get_residual_block<form>(utils::blocks::gravity_field.poisson_term);

        const std::array<int, form::value_block_count> value_sizes{
            f.densityFes->GetTrueVSize(), f.displacementFes->GetTrueVSize(), f.gravityFluxFes->GetTrueVSize(),
            f.gravityPotentialFes->GetTrueVSize()
        };

        const std::array<int, form::residual_block_count> residual_sizes{
            f.gravityFluxFes->GetTrueVSize(), f.gravityPotentialFes->GetTrueVSize()
        };

        const utils::blocks::form_layout<form> layout(value_sizes, residual_sizes);

        mfem::Vector density_true;
        mfem::Vector displacement_true;

        grid_function_to_true_dofs(*f.densityFes, rho, density_true);
        grid_function_to_true_dofs(*f.displacementFes, displacement, displacement_true);

        operators::context::gravity_field::GravityFieldLinearizationContext linearization_context(
            f, *f.domainMapperStateless
        );

        operators::GravityFieldJacobianOperator gravity_jacobian(
            f, *f.domainMapperStateless, linearization_context, layout.value_offsets(), layout.residual_offsets()
        );

        operators::GravityFieldOperator gravity_operator(
            f, *f.domainMapperStateless, linearization_context, layout.value_offsets(), gravity_jacobian
        );

        operators::context::gravity_field::GravityFieldGeometryContext reduced_geometry_context(
            f, *f.domainMapperStateless
        );

        operators::ReducedGravityFieldOperator reduced_operator(
            gravity_operator, reduced_geometry_context, displacement_true
        );

        mfem::Vector right_hand_side;
        reduced_operator.BuildRightHandSide(density_true, right_hand_side);

        MFEM_VERIFY(
            right_hand_side.Size() == reduced_operator.Height(),
            "The reduced gravity right-hand side has the wrong size."
        );

        mfem::BlockVector gravity_state(reduced_operator.GetGravityTrueOffsets());
        gravity_state = 0.0;

        mfem::MINRESSolver minres(f.mesh->GetComm());
        minres.SetOperator(reduced_operator);
        minres.SetPreconditioner(*f.gravityContext.block_prec);
        minres.SetRelTol(args.p.rtol);
        minres.SetAbsTol(args.p.atol);
        minres.SetMaxIter(args.p.max_iters);
        minres.SetPrintLevel(1);
        minres.Mult(right_hand_side, gravity_state);

        MFEM_VERIFY(minres.GetConverged(), "The reduced gravity solve failed to converge.");

        GravitySolution solution(f);

        solution.gradPhi.SetFromTrueDofs(gravity_state.GetBlock(gravity_gradient_residual_block));

        solution.phi.SetFromTrueDofs(gravity_state.GetBlock(gravity_poisson_residual_block));

        return solution;
    }
} // namespace mean_field::physics