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

module mean_field;
import :operators.context.gravity_field;

namespace {
    void validate_displacement(
        const mean_field::fem::FEM &f,
        const mfem::Vector &displacement_true
    ) {
        MFEM_VERIFY(
            f.displacementFes != nullptr,
            "GravityFieldGeometryContext requires the "
            "displacement finite-element space."
        );
        MFEM_VERIFY(
            displacement_true.Size() == f.displacementFes->GetTrueVSize(),
            "GravityFieldGeometryContext received a displacement vector with "
            "the "
            "wrong size."
        );

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

    void validate_linearization_state(
        const mean_field::fem::FEM &f,
        const mean_field::operators::context::gravity_field::
            GravityFieldStateView &state
    ) {
        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."
        );

        MFEM_VERIFY(
            state.density.Size() == f.densityFes->GetTrueVSize(),
            "GravityFieldLinearizationContext received a density vector with "
            "the "
            "wrong size."
        );
        MFEM_VERIFY(
            state.displacement.Size() == f.displacementFes->GetTrueVSize(),
            "GravityFieldLinearizationContext received a displacement vector "
            "with "
            "the wrong size."
        );
        MFEM_VERIFY(
            state.gravity_gradient.Size() == f.gravityFluxFes->GetTrueVSize(),
            "GravityFieldLinearizationContext received a gravity-gradient "
            "vector "
            "with the wrong size."
        );
        MFEM_VERIFY(
            state.gravity_potential.Size() ==
                f.gravityPotentialFes->GetTrueVSize(),
            "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::DomainMapperStateless &domain_mapper
    )
        : m_fem(f),
          m_domain_mapper(domain_mapper) {
        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_true,
        const DiscretizationRevision discretization_revision,
        const DisplacementRevision displacement_revision
    ) {
        validate_displacement(m_fem, displacement_true);

        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;

        GravityFieldGeometryPreparation preparation;

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

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

            mass_operator->Prepare(displacement_true);
            source_operator->Prepare(displacement_true);

            m_mass_operator                     = std::move(mass_operator);
            m_source_operator                   = std::move(source_operator);

            preparation.reconstructed_operators = true;
            preparation.rebuilt_mass_operator   = true;
            preparation.rebuilt_source_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."
            );

            m_mass_operator->Prepare(displacement_true);
            m_source_operator->Prepare(displacement_true);

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

        m_displacement_true                   = displacement_true;
        m_discretization_revision             = discretization_revision;
        m_displacement_revision               = displacement_revision;
        m_is_prepared                         = true;

        preparation.refreshed_variation_state = true;

        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::Vector &GravityFieldGeometryContext::GetDisplacement() const {
        MFEM_VERIFY(
            m_is_prepared,
            "GravityFieldGeometryContext must be prepared before "
            "accessing its displacement."
        );
        return m_displacement_true;
    }

    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::DomainMapperStateless &domain_mapper
    )
        : m_fem(f),
          m_geometry_context(
              f,
              domain_mapper
          ) {
        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_fem, 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         = state.density;
            report.updated_density = true;
        }

        if (gravity_gradient_changed) {
            m_gravity_gradient_true         = state.gravity_gradient;
            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::GetDensity() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its density."
        );
        return m_density_true;
    }

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

    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