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

module mean_field;
import :operators.prepared_hdiv_mass;

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

    int get_operator_size(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "PreparedMappedHDivMassOperator requires the "
                                         "gravity-gradient finite-element space."
        );
        return mean_field::field::make_field_dof_map<mean_field::field::Gravity, DomainSchema>(*f.gravityFluxFes)
            .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;
        }
    }

    int find_representative_element(
        const mean_field::fem::FEM &f,
        const mfem::Array<int> &marker
    ) {
        for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
            const int attribute = f.mesh->GetAttribute(element_id);

            if (attribute > 0 && attribute <= marker.Size() && marker[attribute - 1] != 0) {
                return element_id;
            }
        }

        return -1;
    }

    void validate_uniform_domain_discretization(
        const mean_field::fem::FEM &f,
        const mfem::Array<int> &marker,
        const int representative_element_id
    ) {
        const mfem::FiniteElement &representative_element = *f.gravityFluxFes->GetFE(representative_element_id);
        const mfem::ElementTransformation &representative_transformation =
            *f.mesh->GetElementTransformation(representative_element_id);

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

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

            const mfem::FiniteElement &element                = *f.gravityFluxFes->GetFE(element_id);
            const mfem::ElementTransformation &transformation = *f.mesh->GetElementTransformation(element_id);

            MFEM_VERIFY(
                element.GetGeomType() == representative_element.GetGeomType(),
                "Prepared H(div) mass domains currently require a uniform "
                "element "
                "geometry."
            );
            MFEM_VERIFY(
                element.GetOrder() == representative_element.GetOrder(),
                "Prepared H(div) mass domains currently require a uniform "
                "finite-element order."
            );
            MFEM_VERIFY(
                transformation.OrderW() == representative_transformation.OrderW(),
                "Prepared H(div) mass domains currently require a uniform "
                "geometry-weight order."
            );
        }
    }

    class FrozenMappedHDivMassCoefficient final : public mfem::MatrixCoefficient {
    public:
        FrozenMappedHDivMassCoefficient(
            const mean_field::fem::FEM &f,
            const mean_field::mapping::DomainMapperStateless &domain_mapper,
            const mfem::Vector &displacement_true,
            bool elevates_vacuum
        )
            : MatrixCoefficient(domain_mapper.GetDimension()),
              m_fem(f),
              m_domain_mapper(domain_mapper),
              m_workspace(domain_mapper.GetDimension()),
              m_elevates_vacuum(elevates_vacuum) {
            true_to_local(*m_fem.displacementFes, displacement_true, m_displacement_local);
        }

        void Eval(
            mfem::DenseMatrix &mass_tensor,
            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 H(div) mass coefficient received an invalid element ID."
            );

            const bool element_is_vacuum = transformation.Attribute == m_domain_mapper.GetVacuumElementAttribute();

            if (element_is_vacuum != m_elevates_vacuum) {
                mass_tensor.SetSize(m_domain_mapper.GetDimension());
                mass_tensor = 0.0;
                return;
            }

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

            MFEM_VERIFY(
                status == mean_field::mapping::MappingStatus::valid,
                "Stateless domain mapping failed while preparing the H(div) "
                "mass "
                "operator. Mapping status = "
                    << static_cast<int>(status) << ", element ID = " << element_id
                    << ", element attribute = " << transformation.Attribute
                    << ", coefficient domain = " << (m_elevates_vacuum ? "vacuum" : "stellar")
            );

            const mfem::DenseMatrix &mapping_jacobian = mapping_context.mapping.mapping_jacobian;
            const double mapping_determinant          = mapping_context.mapping.mapping_determinant;

            MFEM_VERIFY(
                std::isfinite(mapping_determinant) && mapping_determinant > 0.0,
                "Prepared H(div) mass operator encountered a non-positive or "
                "non-finite mapping determinant."
            );

            mfem::MultAtB(mapping_jacobian, mapping_jacobian, mass_tensor);
            mass_tensor *= 1.0 / mapping_determinant;
        }

    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::DomainMapperStateless &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::DomainMapperStateless::Workspace m_workspace;
        int m_cached_element_id{-1};
        bool m_elevates_vacuum;
    };
} // namespace

namespace mean_field::operators {
    PreparedMappedHDivMassOperator::PreparedMappedHDivMassOperator(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domain_mapper
    )
        : Operator(get_operator_size(f)),
          m_fem(f),
          m_domain_mapper(domain_mapper),
          m_flux_map(
              field::make_field_dof_map<
                  field::Gravity,
                  DomainSchema>(*f.gravityFluxFes)
          ),
          m_displacement_map(
              field::make_field_dof_map<
                  field::Displacement,
                  DomainSchema>(*f.displacementFes)
          ) {
        MFEM_VERIFY(f.mesh != nullptr, "PreparedMappedHDivMassOperator requires a mesh.");
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "PreparedMappedHDivMassOperator requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "PreparedMappedHDivMassOperator requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "PreparedMappedHDivMassOperator requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "PreparedMappedHDivMassOperator requires the compactification "
                                                     "coordinate."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "PreparedMappedHDivMassOperator requires the quadrature-rule "
                                            "factory."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "The stateless domain-mapper dimension does not match the mesh "
            "dimension."
        );

        utils::populate_element_mask(f.mesh.get(), utils::DOMAINS::STELLAR, m_stellar_marker);
        utils::populate_element_mask(f.mesh.get(), utils::DOMAINS::VACUUM, m_vacuum_marker);

        const int stellar_element_id = find_representative_element(f, m_stellar_marker);
        const int vacuum_element_id  = find_representative_element(f, m_vacuum_marker);

        MFEM_VERIFY(
            stellar_element_id >= 0, "PreparedMappedHDivMassOperator requires "
                                     "at least one stellar element."
        );
        MFEM_VERIFY(
            vacuum_element_id >= 0, "PreparedMappedHDivMassOperator requires at "
                                    "least one compactified vacuum element."
        );

        validate_uniform_domain_discretization(f, m_stellar_marker, stellar_element_id);
        validate_uniform_domain_discretization(f, m_vacuum_marker, vacuum_element_id);
    }

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

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

        m_displacement_true.SetSize(m_displacement_map.full_size());
        m_displacement_map.scatter(displacement, m_displacement_true);

        const int stellar_element_id                        = find_representative_element(m_fem, m_stellar_marker);
        const int vacuum_element_id                         = find_representative_element(m_fem, m_vacuum_marker);

        const mfem::FiniteElement &stellar_element          = *m_fem.gravityFluxFes->GetFE(stellar_element_id);
        const mfem::FiniteElement &vacuum_element           = *m_fem.gravityFluxFes->GetFE(vacuum_element_id);

        mfem::ElementTransformation &stellar_transformation = *m_fem.mesh->GetElementTransformation(stellar_element_id);
        mfem::ElementTransformation &vacuum_transformation  = *m_fem.mesh->GetElementTransformation(vacuum_element_id);

        m_mass_form.reset();
        m_stellar_mass_coefficient.reset();
        m_vacuum_mass_coefficient.reset();

        m_stellar_mass_coefficient =
            std::make_unique<FrozenMappedHDivMassCoefficient>(m_fem, m_domain_mapper, m_displacement_true, false);
        m_vacuum_mass_coefficient =
            std::make_unique<FrozenMappedHDivMassCoefficient>(m_fem, m_domain_mapper, m_displacement_true, true);

        m_mass_form = std::make_unique<mfem::ParBilinearForm>(m_fem.gravityFluxFes.get());
        m_mass_form->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);

        auto stellar_integrator = std::make_unique<mfem::VectorFEMassIntegrator>(*m_stellar_mass_coefficient);
        auto vacuum_integrator  = std::make_unique<mfem::VectorFEMassIntegrator>(*m_vacuum_mass_coefficient);

        m_fem.quadratureFactory->configure_gravity_hdiv_mass(
            *stellar_integrator, quadrature::QuadratureRole::discretization, stellar_element, stellar_transformation,
            utils::DOMAINS::STELLAR, quadrature::MappingKind::general
        );

        m_fem.quadratureFactory->configure_gravity_hdiv_mass(
            *vacuum_integrator, quadrature::QuadratureRole::discretization, vacuum_element, vacuum_transformation,
            utils::DOMAINS::VACUUM, quadrature::MappingKind::kelvin
        );

        m_mass_form->AddDomainIntegrator(stellar_integrator.release(), m_stellar_marker);
        m_mass_form->AddDomainIntegrator(vacuum_integrator.release(), m_vacuum_marker);
        m_mass_form->Assemble();

        m_is_prepared = true;
        ++m_preparation_count;
    }

    void PreparedMappedHDivMassOperator::Mult(
        const mfem::Vector &gravity_gradient,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_is_prepared, "PreparedMappedHDivMassOperator must be prepared "
                           "before Mult is called."
        );
        MFEM_VERIFY(
            m_mass_form != nullptr, "PreparedMappedHDivMassOperator has no "
                                    "assembled partial-assembly form."
        );
        MFEM_VERIFY(
            gravity_gradient.Size() == Width(), "PreparedMappedHDivMassOperator received a gravity-gradient vector "
                                                "with the wrong size."
        );

        m_flux_true.SetSize(m_flux_map.full_size());
        m_action_true.SetSize(m_flux_map.full_size());
        m_flux_map.scatter(gravity_gradient, m_flux_true);
        m_mass_form->Mult(m_flux_true, m_action_true);
        action.SetSize(Height());
        m_flux_map.gather(m_action_true, action);
    }

    bool PreparedMappedHDivMassOperator::IsPrepared() const noexcept {
        return m_is_prepared;
    }

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

    const field::FieldDofMap &PreparedMappedHDivMassOperator::GetFluxMap() const noexcept {
        return m_flux_map;
    }

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