module;
#include <cmath>
#include <cstdint>
#include <memory>
#include <mfem.hpp>
#include <numbers>

module mean_field;
import :operators.prepared_gravity_source;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

    int get_operator_height(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "PreparedMappedGravitySourceOperator requires the "
                                              "gravity-potential "
                                              "finite-element space."
        );
        return mean_field::field::make_field_dof_map<mean_field::field::Gravity, DomainSchema>(*f.gravityPotentialFes)
            .reduced_size();
    }

    int get_operator_width(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.densityFes != nullptr, "PreparedMappedGravitySourceOperator requires the density "
                                     "finite-element space."
        );
        return mean_field::field::make_field_dof_map<mean_field::field::Density, DomainSchema>(*f.densityFes)
            .reduced_size();
    }

    void true_to_local(
        const mfem::ParFiniteElementSpace &finite_element_space,
        const mfem::Vector &true_vector,
        mfem::Vector &local_vector
    ) {
        local_vector.SetSize(finite_element_space.GetVSize());

        const mfem::Operator *prolongation = finite_element_space.GetProlongationMatrix();

        if (prolongation != nullptr) {
            prolongation->Mult(true_vector, local_vector);
        } else {
            local_vector = true_vector;
        }
    }

    void local_to_true(
        const mfem::ParFiniteElementSpace &finite_element_space,
        const mfem::Vector &local_vector,
        mfem::Vector &true_vector
    ) {
        MFEM_VERIFY(local_vector.Size() == finite_element_space.GetVSize(), "Local vector has the wrong size.");

        true_vector.SetSize(finite_element_space.GetTrueVSize());
        true_vector                        = 0.0;

        const mfem::Operator *prolongation = finite_element_space.GetProlongationMatrix();

        if (prolongation != nullptr) {
            prolongation->MultTranspose(local_vector, true_vector);
        } else {
            true_vector = local_vector;
        }
    }

    const mfem::IntegrationRule &get_source_rule(
        const mean_field::fem::FEM &f,
        const mfem::FiniteElement &density_element,
        const mfem::FiniteElement &potential_element,
        const mfem::ElementTransformation &transformation
    ) {
        using GravityField = mean_field::field::Field<mean_field::field::Gravity>;
        MFEM_VERIFY(
            density_element.GetOrder() == mean_field::field::Density::Scalar::familyOrder,
            "The prepared source trial element does not match the registered "
            "density field."
        );
        MFEM_VERIFY(
            potential_element.GetOrder() == mean_field::field::Gravity::Potential::familyOrder,
            "The prepared source test element does not match the registered "
            "gravity potential."
        );
        const mean_field::quadrature::Query query =
            GravityField::make_query<mean_field::field::Gravity::Form::SourceProjection>(
                mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
                mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
            );

        return *f.quadratureFactory->get(query, transformation.GetGeometryType()).integration_rule;
    }

    class FrozenMappedGravitySourceCoefficient final : public mfem::Coefficient {
    public:
        FrozenMappedGravitySourceCoefficient(
            const mean_field::fem::FEM &f,
            const mean_field::mapping::DomainMapper &domain_mapper,
            const mfem::Vector &displacement_true
        )
            : m_fem(f),
              m_domain_mapper(domain_mapper),
              m_workspace(domain_mapper.GetDimension()) {
            true_to_local(*m_fem.displacementFes, displacement_true, m_displacement_local);
        }

        double Eval(
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point
        ) override {
            transformation.SetIntPoint(&integration_point);

            const int element_id = transformation.ElementNo;
            MFEM_VERIFY(
                element_id >= 0 && element_id < m_fem.mesh->GetNE(),
                "Mapped gravity source coefficient received an invalid element "
                "ID."
            );
            if (DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Vacuum>(
                    transformation.Attribute
                )) {
                return 0.0;
            }

            LoadElement(element_id);
            const mean_field::mapping::ElementMappingData mapping_data{
                .displacement = *m_displacement_data, .compactification = *m_compactification_data
            };

            mean_field::mapping::VolumeMappingContext mapping_context;

            const mean_field::mapping::MappingStatus status = m_domain_mapper.EvaluateVolume(
                mapping_data, transformation, integration_point, m_workspace, mapping_context
            );

            if (status != mean_field::mapping::MappingStatus::valid) {
                const mfem::FiniteElement &displacement_element     = *m_fem.displacementFes->GetFE(element_id);
                const mfem::FiniteElement &compactification_element = *m_fem.compactificationFes->GetFE(element_id);

                mfem::Vector displacement_shape(displacement_element.GetDof());
                mfem::Vector compactification_shape(compactification_element.GetDof());
                mfem::Vector reference_position(m_domain_mapper.GetDimension());
                mfem::Vector displacement_value(m_domain_mapper.GetDimension());

                displacement_element.CalcShape(integration_point, displacement_shape);
                compactification_element.CalcShape(integration_point, compactification_shape);
                transformation.Transform(integration_point, reference_position);
                m_displacement_data->GetDofMatrix().MultTranspose(displacement_shape, displacement_value);

                const double compactification_coordinate = m_compactification_data->GetDofs() * compactification_shape;

                MFEM_ABORT(
                    "Stateless domain mapping failed while preparing the "
                    "gravity "
                    "source operator."
                    << "\nMapping status = " << static_cast<int>(status) << "\nElement ID = " << element_id
                    << "\nElement attribute = " << transformation.Attribute
                    << "\nIntegration-point index = " << integration_point.index << "\nIntegration point = <"
                    << integration_point.x << ", " << integration_point.y << ", " << integration_point.z << ">"
                    << "\nReference position = <" << reference_position(0) << ", " << reference_position(1) << ", "
                    << reference_position(2) << ">"
                    << "\nReference radius = " << reference_position.Norml2() << "\nDisplacement value = <"
                    << displacement_value(0) << ", " << displacement_value(1) << ", " << displacement_value(2) << ">"
                    << "\nDisplacement magnitude = " << displacement_value.Norml2()
                    << "\nCompactification coordinate = " << compactification_coordinate
                    << "\nDisplacement ordering = " << static_cast<int>(m_fem.displacementFes->GetOrdering())
                );
            }
            const double mapping_determinant = mapping_context.mapping.mapping_determinant;
            MFEM_VERIFY(
                std::isfinite(mapping_determinant) && mapping_determinant > 0.0,
                "Prepared gravity source operator encountered a non-positive "
                "or "
                "non-finite mapping determinant."
            );

            m_inverse_element_jacobian = mapping_context.quadrature.J_inv;

            return 4.0 * std::numbers::pi * mean_field::utils::G * mapping_determinant;
        }

        [[nodiscard]] const mfem::DenseMatrix &GetInverseElementJacobian() const noexcept {
            return m_inverse_element_jacobian;
        }

    private:
        void LoadElement(const int element_id) {
            if (element_id == m_cached_element_id) {
                return;
            }

            const mfem::FiniteElement &displacement_element     = *m_fem.displacementFes->GetFE(element_id);
            const mfem::FiniteElement &compactification_element = *m_fem.compactificationFes->GetFE(element_id);

            mfem::DofTransformation *displacement_dof_transformation =
                m_fem.displacementFes->GetElementVDofs(element_id, m_displacement_dofs);
            mfem::DofTransformation *compactification_dof_transformation =
                m_fem.compactificationFes->GetElementDofs(element_id, m_compactification_dofs);

            m_displacement_local.GetSubVector(m_displacement_dofs, m_element_displacement);
            m_fem.compactificationCoordinate->GetSubVector(m_compactification_dofs, m_element_compactification);

            if (displacement_dof_transformation != nullptr) {
                displacement_dof_transformation->InvTransformPrimal(m_element_displacement);
            }

            if (compactification_dof_transformation != nullptr) {
                compactification_dof_transformation->InvTransformPrimal(m_element_compactification);
            }

            m_displacement_data = std::make_unique<mean_field::mapping::ElementDisplacementData>(
                mean_field::mapping::ElementDisplacementDataFromElementVDofs(
                    displacement_element, m_element_displacement
                )
            );

            m_compactification_data = std::make_unique<mean_field::mapping::ElementCompactificationData>(
                compactification_element, m_element_compactification
            );

            m_cached_element_id = element_id;
        }

        const mean_field::fem::FEM &m_fem;
        const mean_field::mapping::DomainMapper &m_domain_mapper;

        mfem::Vector m_displacement_local;

        mfem::Array<int> m_displacement_dofs;
        mfem::Array<int> m_compactification_dofs;

        mfem::Vector m_element_displacement;
        mfem::Vector m_element_compactification;

        std::unique_ptr<mean_field::mapping::ElementDisplacementData> m_displacement_data;
        std::unique_ptr<mean_field::mapping::ElementCompactificationData> m_compactification_data;

        mean_field::mapping::DomainMapper::Workspace m_workspace;
        mfem::DenseMatrix m_inverse_element_jacobian;
        int m_cached_element_id{-1};
    };
} // namespace

namespace mean_field::operators {
    PreparedMappedGravitySourceOperator::PreparedMappedGravitySourceOperator(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper
    )
        : Operator(
              get_operator_height(f),
              get_operator_width(f)
          ),
          m_fem(f),
          m_domain_mapper(domain_mapper),
          m_density_map(
              field::make_field_dof_map<
                  field::Density,
                  DomainSchema>(*f.densityFes)
          ),
          m_potential_map(
              field::make_field_dof_map<
                  field::Gravity,
                  DomainSchema>(*f.gravityPotentialFes)
          ),
          m_displacement_map(
              field::make_field_dof_map<
                  field::Displacement,
                  DomainSchema>(*f.displacementFes)
          ) {
        MFEM_VERIFY(f.mesh != nullptr, "PreparedMappedGravitySourceOperator requires a mesh.");
        MFEM_VERIFY(
            f.densityFes != nullptr, "PreparedMappedGravitySourceOperator requires the density "
                                     "finite-element space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "PreparedMappedGravitySourceOperator requires the "
                                              "gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "PreparedMappedGravitySourceOperator requires "
                                          "the displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "PreparedMappedGravitySourceOperator requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr,
            "PreparedMappedGravitySourceOperator requires the compactification "
            "coordinate."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "PreparedMappedGravitySourceOperator "
                                            "requires the quadrature-rule factory."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "The stateless domain-mapper dimension does not match the mesh "
            "dimension."
        );

        m_stellar_marker = utils::domain::make_attribute_marker<utils::domain::Stellar, DomainSchema>(*f.mesh);
    }

    void PreparedMappedGravitySourceOperator::Prepare(const mfem::Vector &displacement) {
        MFEM_VERIFY(
            displacement.Size() == m_displacement_map.reduced_size(),
            "PreparedMappedGravitySourceOperator received a displacement "
            "vector "
            "with the wrong size."
        );

        for (int i = 0; i < displacement.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(displacement(i)), "PreparedMappedGravitySourceOperator received a non-finite "
                                                "displacement value."
            );
        }

        m_is_prepared = false;
        m_displacement_true.SetSize(m_displacement_map.full_size());
        m_displacement_map.scatter(displacement, m_displacement_true);
        m_elements.clear();
        m_elements.reserve(m_fem.mesh->GetNE());

        FrozenMappedGravitySourceCoefficient source_coefficient(m_fem, m_domain_mapper, m_displacement_true);

        for (int element_id = 0; element_id < m_fem.mesh->GetNE(); ++element_id) {
            const int attribute = m_fem.mesh->GetAttribute(element_id);

            if (attribute <= 0 || attribute > m_stellar_marker.Size() || m_stellar_marker[attribute - 1] == 0) {
                continue;
            }

            m_elements.emplace_back();
            ElementPAData &data             = m_elements.back();

            data.element_id                 = element_id;

            data.density_dof_transformation = m_fem.densityFes->GetElementDofs(element_id, data.density_dofs);

            data.potential_dof_transformation =
                m_fem.gravityPotentialFes->GetElementDofs(element_id, data.potential_dofs);

            data.displacement_dof_transformation =
                m_fem.displacementFes->GetElementVDofs(element_id, data.displacement_dofs);

            const mfem::FiniteElement &density_element   = *m_fem.densityFes->GetFE(element_id);

            const mfem::FiniteElement &potential_element = *m_fem.gravityPotentialFes->GetFE(element_id);

            mfem::ElementTransformation &transformation  = *m_fem.mesh->GetElementTransformation(element_id);

            const mfem::IntegrationRule &integration_rule =
                get_source_rule(m_fem, density_element, potential_element, transformation);
            data.integration_rule            = &integration_rule;

            const int quadrature_point_count = integration_rule.GetNPoints();

            const int density_dof_count      = density_element.GetDof();

            const int potential_dof_count    = potential_element.GetDof();

            data.density_basis.SetSize(quadrature_point_count, density_dof_count);

            data.potential_basis.SetSize(quadrature_point_count, potential_dof_count);

            const int dimension = m_fem.mesh->Dimension();
            data.inverse_element_jacobians.SetSize(quadrature_point_count, dimension * dimension);

            data.quadrature_data.SetSize(quadrature_point_count);

            mfem::Vector density_shape(density_dof_count);
            mfem::Vector potential_shape(potential_dof_count);

            for (int quadrature_point = 0; quadrature_point < quadrature_point_count; ++quadrature_point) {
                const mfem::IntegrationPoint &integration_point = integration_rule.IntPoint(quadrature_point);

                transformation.SetIntPoint(&integration_point);

                // CalcPhysShape matches the scalar mixed-mass discretization,
                // including the finite-element map type.
                density_element.CalcPhysShape(transformation, density_shape);

                potential_element.CalcPhysShape(transformation, potential_shape);

                for (int i = 0; i < density_dof_count; ++i) {
                    data.density_basis(quadrature_point, i) = density_shape(i);
                }

                for (int i = 0; i < potential_dof_count; ++i) {
                    data.potential_basis(quadrature_point, i) = potential_shape(i);
                }

                const double coefficient_value = source_coefficient.Eval(transformation, integration_point);

                const mfem::DenseMatrix &inverse_element_jacobian = source_coefficient.GetInverseElementJacobian();
                for (int row = 0; row < dimension; ++row) {
                    for (int column = 0; column < dimension; ++column) {
                        data.inverse_element_jacobians(quadrature_point, row * dimension + column) =
                            inverse_element_jacobian(row, column);
                    }
                }

                transformation.SetIntPoint(&integration_point);

                const double quadrature_value = integration_point.weight * transformation.Weight() * coefficient_value;

                MFEM_VERIFY(
                    std::isfinite(quadrature_value) && quadrature_value > 0.0,
                    "Prepared gravity source operator encountered invalid "
                    "quadrature data on element "
                        << element_id << ", quadrature point " << quadrature_point << "."
                );

                data.quadrature_data(quadrature_point) = quadrature_value;
            }
        }

        MFEM_VERIFY(!m_elements.empty(), "PreparedMappedGravitySourceOperator found no stellar elements.");

        m_is_prepared = true;
        ++m_preparation_count;
    }
    void PreparedMappedGravitySourceOperator::Mult(
        const mfem::Vector &density,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedGravitySourceOperator must be prepared before "
                           "Mult is called."
        );

        MFEM_VERIFY(
            density.Size() == Width(), "PreparedMappedGravitySourceOperator received a density vector "
                                       "with the wrong size."
        );

        m_density_true.SetSize(m_density_map.full_size());
        m_density_map.scatter(density, m_density_true);

        mfem::Vector density_local;

        true_to_local(*m_fem.densityFes, m_density_true, density_local);

        mfem::Vector local_action(m_fem.gravityPotentialFes->GetVSize());
        local_action = 0.0;

        mfem::Vector element_density;
        mfem::Vector quadrature_density;
        mfem::Vector element_action;

        for (const ElementPAData &data : m_elements) {
            density_local.GetSubVector(data.density_dofs, element_density);

            if (data.density_dof_transformation != nullptr) {
                data.density_dof_transformation->InvTransformPrimal(element_density);
            }

            quadrature_density.SetSize(data.quadrature_data.Size());

            // B_density * x_e
            data.density_basis.Mult(element_density, quadrature_density);

            // D * B_density * x_e
            for (int q = 0; q < quadrature_density.Size(); ++q) {
                quadrature_density(q) *= data.quadrature_data(q);
            }

            element_action.SetSize(data.potential_dofs.Size());

            // B_potential^T * D * B_density * x_e
            data.potential_basis.MultTranspose(quadrature_density, element_action);

            if (data.potential_dof_transformation != nullptr) {
                data.potential_dof_transformation->TransformDual(element_action);
            }

            local_action.AddElementVector(data.potential_dofs, element_action);
        }

        local_to_true(*m_fem.gravityPotentialFes, local_action, m_action_true);
        action.SetSize(Height());
        m_potential_map.gather(m_action_true, action);
    }

    void PreparedMappedGravitySourceOperator::MultDisplacementVariationTrue(
        const mfem::Vector &densityTrue,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &actionVariationTrue
    ) const {
        MFEM_VERIFY(
            m_is_prepared,
            "PreparedMappedGravitySourceOperator must be prepared before applying a displacement variation."
        );
        MFEM_VERIFY(
            densityTrue.Size() == m_fem.densityFes->GetTrueVSize(), "The full density vector has the wrong size."
        );
        MFEM_VERIFY(
            displacementVariationTrue.Size() == m_fem.displacementFes->GetTrueVSize(),
            "The full displacement variation has the wrong size."
        );

        true_to_local(*m_fem.densityFes, densityTrue, m_density_local);
        true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacement_variation_local);

        m_local_variation_action.SetSize(m_fem.gravityPotentialFes->GetVSize());
        m_local_variation_action = 0.0;

        const int dimension      = m_fem.mesh->Dimension();

        for (const ElementPAData &data : m_elements) {
            MFEM_VERIFY(
                data.integration_rule != nullptr,
                "Prepared gravity source displacement variation has no integration rule."
            );

            m_density_local.GetSubVector(data.density_dofs, m_element_density);
            m_displacement_variation_local.GetSubVector(data.displacement_dofs, m_element_displacement_variation);

            if (data.density_dof_transformation != nullptr) {
                data.density_dof_transformation->InvTransformPrimal(m_element_density);
            }
            if (data.displacement_dof_transformation != nullptr) {
                data.displacement_dof_transformation->InvTransformPrimal(m_element_displacement_variation);
            }

            const mfem::FiniteElement &displacement_element       = *m_fem.displacementFes->GetFE(data.element_id);
            const mapping::ElementDisplacementData direction_data = mapping::ElementDisplacementDataFromElementVDofs(
                displacement_element, m_element_displacement_variation
            );
            const mfem::DenseMatrix &direction_dofs = direction_data.GetDofMatrix();

            MFEM_VERIFY(
                data.inverse_element_jacobians.Height() == data.integration_rule->GetNPoints() &&
                    data.inverse_element_jacobians.Width() == dimension * dimension,
                "Prepared gravity source inverse-Jacobian data has an incompatible size."
            );

            m_reference_displacement_dshape.SetSize(displacement_element.GetDof(), dimension);
            m_reference_displacement_jacobian.SetSize(dimension, dimension);
            m_quadrature_variation_action.SetSize(data.integration_rule->GetNPoints());
            data.density_basis.Mult(m_element_density, m_quadrature_variation_action);

            for (int quadrature_point = 0; quadrature_point < data.integration_rule->GetNPoints(); ++quadrature_point) {
                const mfem::IntegrationPoint &integration_point = data.integration_rule->IntPoint(quadrature_point);
                displacement_element.CalcDShape(integration_point, m_reference_displacement_dshape);
                mfem::MultAtB(direction_dofs, m_reference_displacement_dshape, m_reference_displacement_jacobian);

                double logarithmic_jacobian_variation{0.0};
                for (int row = 0; row < dimension; ++row) {
                    for (int column = 0; column < dimension; ++column) {
                        logarithmic_jacobian_variation +=
                            data.inverse_element_jacobians(quadrature_point, row * dimension + column) *
                            m_reference_displacement_jacobian(column, row);
                    }
                }

                m_quadrature_variation_action(quadrature_point) *=
                    data.quadrature_data(quadrature_point) * logarithmic_jacobian_variation;
                MFEM_VERIFY(
                    std::isfinite(m_quadrature_variation_action(quadrature_point)),
                    "Prepared gravity source displacement variation encountered a non-finite quadrature value."
                );
            }

            m_element_variation_action.SetSize(data.potential_dofs.Size());
            data.potential_basis.MultTranspose(m_quadrature_variation_action, m_element_variation_action);

            if (data.potential_dof_transformation != nullptr) {
                data.potential_dof_transformation->TransformDual(m_element_variation_action);
            }
            m_local_variation_action.AddElementVector(data.potential_dofs, m_element_variation_action);
        }

        local_to_true(*m_fem.gravityPotentialFes, m_local_variation_action, actionVariationTrue);
    }

    void PreparedMappedGravitySourceOperator::MultTranspose(
        const mfem::Vector &potential,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedGravitySourceOperator must be prepared before "
                           "MultTranspose is called."
        );

        MFEM_VERIFY(
            potential.Size() == Height(), "PreparedMappedGravitySourceOperator received a potential vector "
                                          "with the wrong size."
        );

        m_potential_true.SetSize(m_potential_map.full_size());
        m_potential_map.scatter(potential, m_potential_true);

        mfem::Vector potential_local;

        true_to_local(*m_fem.gravityPotentialFes, m_potential_true, potential_local);

        mfem::Vector local_action(m_fem.densityFes->GetVSize());
        local_action = 0.0;

        mfem::Vector element_potential;
        mfem::Vector quadrature_potential;
        mfem::Vector element_action;

        for (const ElementPAData &data : m_elements) {
            potential_local.GetSubVector(data.potential_dofs, element_potential);

            if (data.potential_dof_transformation != nullptr) {
                data.potential_dof_transformation->InvTransformPrimal(element_potential);
            }

            quadrature_potential.SetSize(data.quadrature_data.Size());

            data.potential_basis.Mult(element_potential, quadrature_potential);

            for (int q = 0; q < quadrature_potential.Size(); ++q) {
                quadrature_potential(q) *= data.quadrature_data(q);
            }

            element_action.SetSize(data.density_dofs.Size());

            data.density_basis.MultTranspose(quadrature_potential, element_action);

            if (data.density_dof_transformation != nullptr) {
                data.density_dof_transformation->TransformDual(element_action);
            }

            local_action.AddElementVector(data.density_dofs, element_action);
        }

        local_to_true(*m_fem.densityFes, local_action, m_action_true);
        action.SetSize(Width());
        m_density_map.gather(m_action_true, action);
    }
    bool PreparedMappedGravitySourceOperator::IsPrepared() const noexcept {
        return m_is_prepared;
    }

    std::uint64_t PreparedMappedGravitySourceOperator::GetPreparationCount() const noexcept {
        return m_preparation_count;
    }

    const field::FieldDofMap &PreparedMappedGravitySourceOperator::GetDensityMap() const noexcept {
        return m_density_map;
    }

    const field::FieldDofMap &PreparedMappedGravitySourceOperator::GetPotentialMap() const noexcept {
        return m_potential_map;
    }

    const field::FieldDofMap &PreparedMappedGravitySourceOperator::GetDisplacementMap() const noexcept {
        return m_displacement_map;
    }
} // namespace mean_field::operators