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

module mean_field;
import :operators.context.gravity_field;

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

    void true_to_local(
        const mfem::ParFiniteElementSpace &finite_element_space,
        const mfem::Vector &true_vector,
        mfem::Vector &local_vector
    ) {
        MFEM_VERIFY(
            true_vector.Size() == finite_element_space.GetTrueVSize(),
            "True-DOF operator received an input vector with the wrong size."
        );

        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(),
            "True-DOF operator produced a local vector with the wrong size."
        );

        true_vector.SetSize(finite_element_space.GetTrueVSize());
        const mfem::Operator *prolongation = finite_element_space.GetProlongationMatrix();
        if (prolongation != nullptr) {
            prolongation->MultTranspose(local_vector, true_vector);
        } else {
            true_vector = local_vector;
        }
    }

    bool communicator_has_single_rank(const MPI_Comm communicator) {
        int size = 0;
        MFEM_VERIFY(MPI_Comm_size(communicator, &size) == MPI_SUCCESS, "Failed to query the MPI communicator size.");
        MFEM_VERIFY(size > 0, "The MPI communicator must contain at least one rank.");
        return size == 1;
    }

    class TrueDofParMixedBilinearFormOperator final : public mfem::Operator {
    public:
        TrueDofParMixedBilinearFormOperator(
            const mfem::ParFiniteElementSpace &trial_space,
            const mfem::ParFiniteElementSpace &test_space,
            std::unique_ptr<mfem::ParMixedBilinearForm> local_form
        )
            : Operator(
                  test_space.GetTrueVSize(),
                  trial_space.GetTrueVSize()
              ),
              m_trial_space(trial_space),
              m_test_space(test_space),
              m_local_form(std::move(local_form)),
              m_single_rank(communicator_has_single_rank(trial_space.GetComm())) {
            int communicators_compare = MPI_UNEQUAL;
            MFEM_VERIFY(
                MPI_Comm_compare(trial_space.GetComm(), test_space.GetComm(), &communicators_compare) == MPI_SUCCESS,
                "Failed to compare mixed-operator MPI communicators."
            );
            MFEM_VERIFY(
                communicators_compare == MPI_IDENT || communicators_compare == MPI_CONGRUENT,
                "True-DOF mixed operator requires congruent trial and test communicators."
            );
            MFEM_VERIFY(m_local_form != nullptr, "True-DOF mixed operator requires a local bilinear form.");
            MFEM_VERIFY(
                m_local_form->Width() == m_trial_space.GetVSize(),
                "True-DOF mixed operator received an incompatible trial space."
            );
            MFEM_VERIFY(
                m_local_form->Height() == m_test_space.GetVSize(),
                "True-DOF mixed operator received an incompatible test space."
            );
        }

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override {
            MFEM_VERIFY(input.Size() == Width(), "True-DOF mixed operator received an input with the wrong size.");

            if (m_single_rank) [[likely]] {
                output.SetSize(Height());
                m_local_form->Mult(input, output);
                return;
            }

            true_to_local(m_trial_space, input, m_trial_local);
            m_test_local.SetSize(m_test_space.GetVSize());
            m_local_form->Mult(m_trial_local, m_test_local);
            local_to_true(m_test_space, m_test_local, output);
        }

        void MultTranspose(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override {
            MFEM_VERIFY(input.Size() == Height(), "True-DOF mixed transpose received an input with the wrong size.");

            if (m_single_rank) [[likely]] {
                output.SetSize(Width());
                m_local_form->MultTranspose(input, output);
                return;
            }

            true_to_local(m_test_space, input, m_test_local);
            m_trial_local.SetSize(m_trial_space.GetVSize());
            m_local_form->MultTranspose(m_test_local, m_trial_local);
            local_to_true(m_trial_space, m_trial_local, output);
        }

    private:
        const mfem::ParFiniteElementSpace &m_trial_space;
        const mfem::ParFiniteElementSpace &m_test_space;
        std::unique_ptr<mfem::ParMixedBilinearForm> m_local_form;
        mutable mfem::Vector m_trial_local;
        mutable mfem::Vector m_test_local;
        bool m_single_rank;
    };

    [[nodiscard]] std::unique_ptr<mfem::Operator> make_divergence_operator(const mean_field::fem::FEM &f) {
        auto divergence =
            std::make_unique<mfem::ParMixedBilinearForm>(f.gravityFluxFes.get(), f.gravityPotentialFes.get());

        divergence->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);

        auto integrator                                   = std::make_unique<mfem::VectorFEDivergenceIntegrator>();

        const mfem::FiniteElement &trialElement           = *f.gravityFluxFes->GetTypicalFE();
        const mfem::FiniteElement &testElement            = *f.gravityPotentialFes->GetTypicalFE();
        const mfem::ElementTransformation &transformation = *f.mesh->GetElementTransformation(0);

        f.quadratureFactory->configure_gravity_divergence(
            *integrator, mean_field::quadrature::QuadratureRole::discretization, trialElement, testElement,
            transformation, mean_field::utils::DOMAINS::ALL, mean_field::quadrature::MappingKind::none
        );

        divergence->AddDomainIntegrator(integrator.release());
        divergence->Assemble();

        return std::make_unique<TrueDofParMixedBilinearFormOperator>(
            *f.gravityFluxFes, *f.gravityPotentialFes, std::move(divergence)
        );
    }

    void validate_displacement(
        const mean_field::field::FieldDofMap &displacement_map,
        const mfem::Vector &displacement
    ) {
        MFEM_VERIFY(
            displacement.Size() == displacement_map.reduced_size(),
            "GravityFieldGeometryContext received a displacement vector with "
            "the "
            "wrong size."
        );

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

    void validate_linearization_state(
        const mean_field::field::FieldDofMap &density_map,
        const mean_field::field::FieldDofMap &displacement_map,
        const mean_field::field::FieldDofMap &gravity_gradient_map,
        const mean_field::field::FieldDofMap &gravity_potential_map,
        const mean_field::operators::context::gravity_field::GravityFieldStateView &state
    ) {
        MFEM_VERIFY(
            state.density.Size() == density_map.reduced_size(),
            "GravityFieldLinearizationContext received a density vector with "
            "the "
            "wrong size."
        );
        MFEM_VERIFY(
            state.displacement.Size() == displacement_map.reduced_size(),
            "GravityFieldLinearizationContext received a displacement vector "
            "with "
            "the wrong size."
        );
        MFEM_VERIFY(
            state.gravity_gradient.Size() == gravity_gradient_map.reduced_size(),
            "GravityFieldLinearizationContext received a gravity-gradient "
            "vector "
            "with the wrong size."
        );
        MFEM_VERIFY(
            state.gravity_potential.Size() == gravity_potential_map.reduced_size(),
            "GravityFieldLinearizationContext received a gravity-potential "
            "vector "
            "with the wrong size."
        );

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

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

        for (int i = 0; i < state.gravity_gradient.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(state.gravity_gradient(i)), "GravityFieldLinearizationContext received a non-finite "
                                                          "gravity-gradient value."
            );
        }

        for (int i = 0; i < state.gravity_potential.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(state.gravity_potential(i)), "GravityFieldLinearizationContext received a non-finite "
                                                           "gravity-potential value."
            );
        }
    }
} // namespace

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

    GravityFieldGeometryPreparation GravityFieldGeometryContext::Prepare(
        const mfem::Vector &displacement,
        const DiscretizationRevision discretization_revision,
        const DisplacementRevision displacement_revision
    ) {
        return PrepareImpl(
            displacement, discretization_revision, displacement_revision, PreparationMode::linearization
        );
    }

    GravityFieldGeometryPreparation GravityFieldGeometryContext::PreparePrimal(
        const mfem::Vector &displacement,
        const DiscretizationRevision discretization_revision,
        const DisplacementRevision displacement_revision
    ) {
        return PrepareImpl(displacement, discretization_revision, displacement_revision, PreparationMode::primal);
    }

    GravityFieldGeometryPreparation GravityFieldGeometryContext::PrepareImpl(
        const mfem::Vector &displacement,
        const DiscretizationRevision discretization_revision,
        const DisplacementRevision displacement_revision,
        const PreparationMode mode
    ) {
        validate_displacement(m_displacement_map, displacement);

        if (m_is_prepared) {
            MFEM_VERIFY(
                discretization_revision >= m_discretization_revision,
                "GravityFieldGeometryContext received an older discretization "
                "revision."
            );
            MFEM_VERIFY(
                displacement_revision >= m_displacement_revision,
                "GravityFieldGeometryContext received an older displacement "
                "revision."
            );
        }

        const bool discretization_changed = !m_is_prepared || discretization_revision != m_discretization_revision;
        const bool displacement_changed   = !m_is_prepared || displacement_revision != m_displacement_revision;
        const bool requires_variation     = mode == PreparationMode::linearization;
        const bool variation_upgrade      = requires_variation && !m_variation_state_prepared;

        GravityFieldGeometryPreparation preparation;

        if (!discretization_changed && !displacement_changed && !variation_upgrade) {
            return preparation;
        }

        const auto prepare_mass = [&](PreparedMappedHDivMassOperator &mass_operator) {
            if (requires_variation) {
                mass_operator.Prepare(displacement);
            } else {
                mass_operator.PreparePrimal(displacement);
            }
        };
        const auto prepare_source = [&](PreparedMappedGravitySourceOperator &source_operator) {
            if (requires_variation) {
                source_operator.Prepare(displacement);
            } else {
                source_operator.PreparePrimal(displacement);
            }
        };

        if (discretization_changed) {
            auto mass_operator       = std::make_unique<PreparedMappedHDivMassOperator>(m_fem, m_domain_mapper);
            auto source_operator     = std::make_unique<PreparedMappedGravitySourceOperator>(m_fem, m_domain_mapper);
            auto divergence_operator = make_divergence_operator(m_fem);
            auto transpose_divergence_operator = std::make_unique<mfem::TransposeOperator>(divergence_operator.get());

            prepare_mass(*mass_operator);
            prepare_source(*source_operator);

            m_mass_operator                         = std::move(mass_operator);
            m_source_operator                       = std::move(source_operator);
            m_divergence_operator                   = std::move(divergence_operator);
            m_transpose_divergence_operator         = std::move(transpose_divergence_operator);

            preparation.reconstructed_operators     = true;
            preparation.rebuilt_mass_operator       = true;
            preparation.rebuilt_source_operator     = true;
            preparation.rebuilt_divergence_operator = true;
        } else {
            MFEM_VERIFY(
                m_mass_operator != nullptr, "GravityFieldGeometryContext has "
                                            "no prepared H(div) mass operator."
            );
            MFEM_VERIFY(
                m_source_operator != nullptr, "GravityFieldGeometryContext has no prepared gravity source "
                                              "operator."
            );

            prepare_mass(*m_mass_operator);
            prepare_source(*m_source_operator);

            preparation.rebuilt_mass_operator   = true;
            preparation.rebuilt_source_operator = true;
        }

        m_displacement_true.SetSize(m_displacement_map.full_size());
        m_displacement_map.scatter(displacement, m_displacement_true);
        m_discretization_revision             = discretization_revision;
        m_displacement_revision               = displacement_revision;
        m_is_prepared                         = true;
        m_variation_state_prepared            = requires_variation;

        preparation.refreshed_variation_state = requires_variation;

        return preparation;
    }

    const PreparedMappedHDivMassOperator &GravityFieldGeometryContext::GetMassOperator() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldGeometryContext must be prepared before "
                           "accessing its mass operator."
        );
        MFEM_VERIFY(m_mass_operator != nullptr, "GravityFieldGeometryContext has no prepared H(div) mass operator.");
        return *m_mass_operator;
    }

    const PreparedMappedGravitySourceOperator &GravityFieldGeometryContext::GetSourceOperator() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldGeometryContext must be prepared before "
                           "accessing its source operator."
        );
        MFEM_VERIFY(
            m_source_operator != nullptr, "GravityFieldGeometryContext has no "
                                          "prepared gravity source operator."
        );
        return *m_source_operator;
    }

    const mfem::Operator &GravityFieldGeometryContext::GetDivergenceOperator() const {
        MFEM_VERIFY(m_is_prepared, "GravityFieldGeometryContext must be prepared before accessing divergence.");
        MFEM_VERIFY(m_divergence_operator != nullptr, "GravityFieldGeometryContext has no divergence operator.");
        return *m_divergence_operator;
    }

    const mfem::Operator &GravityFieldGeometryContext::GetTransposeDivergenceOperator() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldGeometryContext must be prepared before accessing transpose divergence."
        );
        MFEM_VERIFY(
            m_transpose_divergence_operator != nullptr,
            "GravityFieldGeometryContext has no transpose-divergence operator."
        );
        return *m_transpose_divergence_operator;
    }

    const mfem::Vector &GravityFieldGeometryContext::GetDisplacementTrue() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldGeometryContext must be prepared before "
                           "accessing its displacement."
        );
        return m_displacement_true;
    }

    const field::FieldDofMap &GravityFieldGeometryContext::GetDisplacementMap() const noexcept {
        return m_displacement_map;
    }

    DiscretizationRevision GravityFieldGeometryContext::GetDiscretizationRevision() const noexcept {
        return m_discretization_revision;
    }

    DisplacementRevision GravityFieldGeometryContext::GetDisplacementRevision() const noexcept {
        return m_displacement_revision;
    }

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

    GravityFieldLinearizationContext::GravityFieldLinearizationContext(
        const fem::FEM &f,
        const mapping::DomainMapper &domain_mapper
    )
        : m_fem(f),
          m_geometry_context(
              f,
              domain_mapper
          ),
          m_density_map(
              field::make_field_dof_map<
                  field::Density,
                  DomainSchema>(*f.densityFes)
          ),
          m_gravity_gradient_map(
              field::make_field_dof_map<
                  field::Gravity,
                  DomainSchema>(*f.gravityFluxFes)
          ),
          m_gravity_potential_map(
              field::make_field_dof_map<
                  field::Gravity,
                  DomainSchema>(*f.gravityPotentialFes)
          ) {
        MFEM_VERIFY(
            f.densityFes != nullptr, "GravityFieldLinearizationContext "
                                     "requires the density finite-element "
                                     "space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "GravityFieldLinearizationContext requires the gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "GravityFieldLinearizationContext requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldLinearizationContext requires "
                                          "the displacement finite-element space."
        );
    }

    GravityFieldPreparationReport GravityFieldLinearizationContext::Prepare(
        const GravityFieldStateView &state,
        const GravityFieldRevisions &revisions
    ) {
        validate_linearization_state(
            m_density_map, m_geometry_context.GetDisplacementMap(), m_gravity_gradient_map, m_gravity_potential_map,
            state
        );

        if (m_is_prepared) {
            MFEM_VERIFY(
                revisions.discretization >= m_revisions.discretization,
                "GravityFieldLinearizationContext received an older "
                "discretization "
                "revision."
            );
            MFEM_VERIFY(
                revisions.displacement >= m_revisions.displacement,
                "GravityFieldLinearizationContext received an older "
                "displacement "
                "revision."
            );
            MFEM_VERIFY(
                revisions.density >= m_revisions.density, "GravityFieldLinearizationContext received an older density "
                                                          "revision."
            );
            MFEM_VERIFY(
                revisions.gravity_gradient >= m_revisions.gravity_gradient,
                "GravityFieldLinearizationContext received an older "
                "gravity-gradient "
                "revision."
            );
            MFEM_VERIFY(
                revisions.gravity_potential >= m_revisions.gravity_potential,
                "GravityFieldLinearizationContext received an older "
                "gravity-potential revision."
            );
        }

        const bool discretization_changed = !m_is_prepared || revisions.discretization != m_revisions.discretization;
        const bool density_changed =
            !m_is_prepared || discretization_changed || revisions.density != m_revisions.density;
        const bool gravity_gradient_changed =
            !m_is_prepared || discretization_changed || revisions.gravity_gradient != m_revisions.gravity_gradient;

        GravityFieldPreparationReport report;

        report.geometry =
            m_geometry_context.Prepare(state.displacement, revisions.discretization, revisions.displacement);

        if (density_changed) {
            m_density_true.SetSize(m_density_map.full_size());
            m_density_map.scatter(state.density, m_density_true);
            report.updated_density = true;
        }

        if (gravity_gradient_changed) {
            m_gravity_gradient_true.SetSize(m_gravity_gradient_map.full_size());
            m_gravity_gradient_map.scatter(state.gravity_gradient, m_gravity_gradient_true);
            report.updated_gravity_gradient = true;
        }

        m_revisions   = revisions;
        m_is_prepared = true;

        return report;
    }

    const GravityFieldGeometryContext &GravityFieldLinearizationContext::GetGeometryContext() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its geometry context."
        );
        return m_geometry_context;
    }

    const mfem::Vector &GravityFieldLinearizationContext::GetDensityTrue() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its density."
        );
        return m_density_true;
    }

    const mfem::Vector &GravityFieldLinearizationContext::GetGravityGradientTrue() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its gravity gradient."
        );
        return m_gravity_gradient_true;
    }

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

    const field::FieldDofMap &GravityFieldLinearizationContext::GetDisplacementMap() const noexcept {
        return m_geometry_context.GetDisplacementMap();
    }

    const field::FieldDofMap &GravityFieldLinearizationContext::GetGravityGradientMap() const noexcept {
        return m_gravity_gradient_map;
    }

    const field::FieldDofMap &GravityFieldLinearizationContext::GetGravityPotentialMap() const noexcept {
        return m_gravity_potential_map;
    }

    const GravityFieldRevisions &GravityFieldLinearizationContext::GetRevisions() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its revisions."
        );
        return m_revisions;
    }

    bool GravityFieldLinearizationContext::IsPrepared() const noexcept {
        return m_is_prepared;
    }
} // namespace mean_field::operators::context::gravity_field