module;
#include "mfem.hpp"
#include <array>
#include <cmath>

module mean_field;

namespace mean_field::physics {
    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;
        using DomainSchema = utils::domain::CoreEnvelopeVacuumDomainSchema;
        mapping::GridFunctionMappingEvaluator mapping_evaluator(
            *fem.domainMapperStateless, *fem.displacement,
            *fem.compactificationCoordinate
        );

        for (int i = 0; i < fem.mesh->GetNE(); ++i) {
            if (!DomainSchema::template attribute_belongs_to<utils::domain::Stellar>(
                    fem.mesh->GetAttribute(i)))
                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);

                mapping::VolumeMappingContext mapping_context;
                MFEM_VERIFY(
                    mapping_evaluator.EvaluateVolume(*trans, ip, mapping_context) ==
                        mapping::MappingStatus::valid,
                    "Quadrupole integration encountered an invalid mapping."
                );
                const double weight = mapping_context.quadrature.weight;

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

                const mfem::Vector &phys_point = mapping_context.mapping.physical_position;

                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;
    }

    GravitySolution solve_gravity_field(
        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(
            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."
        );
        MFEM_VERIFY(args.p.max_iters > 0, "Gravity solve requires a positive MINRES iteration limit.");

        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);

        using DomainSchema                   = utils::domain::CoreEnvelopeVacuumDomainSchema;
        const field::FieldDofGridFunctionAdapter density_adapter =
            field::make_field_dof_grid_function_adapter<field::Density, DomainSchema>(*f.densityFes);
        const field::FieldDofGridFunctionAdapter displacement_adapter =
            field::make_field_dof_grid_function_adapter<field::Displacement, DomainSchema>(*f.displacementFes);
        const field::FieldDofGridFunctionAdapter gravity_flux_adapter =
            field::make_field_dof_grid_function_adapter<field::Gravity, DomainSchema>(*f.gravityFluxFes);
        const field::FieldDofGridFunctionAdapter gravity_potential_adapter =
            field::make_field_dof_grid_function_adapter<field::Gravity, DomainSchema>(*f.gravityPotentialFes);

        const field::FieldDofMap &density_map           = density_adapter.dof_map();
        const field::FieldDofMap &displacement_map      = displacement_adapter.dof_map();
        const field::FieldDofMap &gravity_flux_map      = gravity_flux_adapter.dof_map();
        const field::FieldDofMap &gravity_potential_map = gravity_potential_adapter.dof_map();

        const std::array<int, form::value_block_count> value_sizes{
            density_map.reduced_size(), displacement_map.reduced_size(), gravity_flux_map.reduced_size(),
            gravity_potential_map.reduced_size()
        };

        const std::array<int, form::residual_block_count> residual_sizes{
            gravity_flux_map.reduced_size(), gravity_potential_map.reduced_size()
        };

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

        const mfem::Vector density              = density_adapter.gather(rho);
        const mfem::Vector reduced_displacement = displacement_adapter.gather(displacement);

        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, reduced_displacement
        );
        operators::ReducedGravityFieldPreconditioner reduced_preconditioner(f, reduced_geometry_context);

        mfem::Vector right_hand_side;
        reduced_operator.BuildRightHandSide(density, 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.GetGravityOffsets());
        gravity_state = 0.0;

        mfem::MINRESSolver minres(f.mesh->GetComm());
        minres.SetOperator(reduced_operator);
        minres.SetPreconditioner(reduced_preconditioner);
        minres.SetRelTol(args.p.rtol);
        minres.SetAbsTol(args.p.atol);
        minres.SetMaxIter(args.p.max_iters);
        // minres.SetPrintLevel(args.verbose ? 1 : 0);
        minres.SetPrintLevel(0);
        minres.Mult(right_hand_side, gravity_state);

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

        GravitySolution solution(f);

        gravity_flux_adapter.scatter(
            gravity_state.GetBlock(gravity_gradient_residual_block), solution.gradPhi
        );
        gravity_potential_adapter.scatter(
            gravity_state.GetBlock(gravity_poisson_residual_block), solution.phi
        );

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