module;

#include <mfem.hpp>

module mean_field;

import :operators.kernels.gravity_displacement_force;
import :operators.prepared_gravity_displacement_force;

namespace {
    [[nodiscard]] bool relevant_revisions_match(
        const mean_field::operators::context::gravity_field::GravityFieldRevisions &left,
        const mean_field::operators::context::gravity_field::GravityFieldRevisions &right
    ) noexcept {
        return left.discretization == right.discretization && left.displacement == right.displacement &&
               left.density == right.density && left.gravity_gradient == right.gravity_gradient;
    }
} // namespace

namespace mean_field::operators {
    PreparedGravityDisplacementForceOperator::PreparedGravityDisplacementForceOperator(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const context::gravity_field::GravityFieldLinearizationContext &gravityContext
    )
        : m_fem(f),
          m_domainMapper(domainMapper),
          m_gravityContext(gravityContext) {
        MFEM_VERIFY(m_fem.mesh != nullptr, "PreparedGravityDisplacementForceOperator requires a mesh.");

        MFEM_VERIFY(
            m_fem.densityFes != nullptr && m_fem.gravityFluxFes != nullptr && m_fem.displacementFes != nullptr,
            "PreparedGravityDisplacementForceOperator requires density, "
            "gravity-gradient, and displacement finite-element spaces."
        );

        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_fem.mesh->Dimension(),
            "PreparedGravityDisplacementForceOperator received a mapper "
            "with the wrong dimension."
        );
    }

    PreparedGravityDisplacementForceReport PreparedGravityDisplacementForceOperator::Prepare() {
        MFEM_VERIFY(
            m_gravityContext.IsPrepared(), "PreparedGravityDisplacementForceOperator requires the shared "
                                           "gravity linearization context to be prepared first."
        );

        const context::gravity_field::GravityFieldRevisions &requestedRevisions = m_gravityContext.GetRevisions();

        if (m_isPrepared && relevant_revisions_match(requestedRevisions, m_preparedRevisions)) {
            return {};
        }

        kernels::apply_gravity_displacement_force_residual(
            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityContext.GetGravityGradientTrue(),
            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        m_cachedResidual.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, m_cachedResidual);

        m_preparedRevisions = requestedRevisions;
        ++m_residualPreparationCount;
        m_isPrepared = true;

        return {.preparedResidual = true};
    }

    void PreparedGravityDisplacementForceOperator::BuildResidual(mfem::Vector &residual) const {
        VerifyPrepared();
        residual = m_cachedResidual;
        ++m_residualApplicationCount;
    }

    void PreparedGravityDisplacementForceOperator::ApplyDensityJacobianAction(
        const mfem::Vector &densityVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);

        kernels::apply_gravity_displacement_force_density_action(
            m_fem, m_domainMapper, m_densityVariationTrue, m_gravityContext.GetGravityGradientTrue(),
            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_densityJacobianStatistics.applications;
    }

    void PreparedGravityDisplacementForceOperator::ApplyGravityGradientJacobianAction(
        const mfem::Vector &gravityGradientVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_gravityGradientVariationTrue.SetSize(m_gravityContext.GetGravityGradientMap().full_size());
        m_gravityContext.GetGravityGradientMap().scatter(gravityGradientVariation, m_gravityGradientVariationTrue);

        kernels::apply_gravity_displacement_force_gradient_action(
            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityGradientVariationTrue,
            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_gravityGradientJacobianStatistics.applications;
    }

    void PreparedGravityDisplacementForceOperator::ApplyDisplacementJacobianAction(
        const mfem::Vector &displacementVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        kernels::apply_gravity_displacement_force_displacement_action(
            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityContext.GetGravityGradientTrue(),
            m_displacementVariationTrue, m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_displacementJacobianStatistics.applications;
    }

    void PreparedGravityDisplacementForceOperator::ApplyCompleteJacobianAction(
        const mfem::Vector &densityVariation,
        const mfem::Vector &displacementVariation,
        const mfem::Vector &gravityGradientVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
        m_gravityGradientVariationTrue.SetSize(m_gravityContext.GetGravityGradientMap().full_size());
        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_gravityContext.GetGravityGradientMap().scatter(gravityGradientVariation, m_gravityGradientVariationTrue);
        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        kernels::apply_gravity_displacement_force_complete_action(
            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_densityVariationTrue,
            m_gravityContext.GetGravityGradientTrue(), m_gravityGradientVariationTrue, m_displacementVariationTrue,
            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_densityJacobianStatistics.applications;
        ++m_gravityGradientJacobianStatistics.applications;
        ++m_displacementJacobianStatistics.applications;
        ++m_completeJacobianStatistics.applications;
    }

    bool PreparedGravityDisplacementForceOperator::IsPrepared() const noexcept {
        if (!m_isPrepared || !m_gravityContext.IsPrepared()) {
            return false;
        }

        return relevant_revisions_match(m_gravityContext.GetRevisions(), m_preparedRevisions);
    }

    std::uint64_t PreparedGravityDisplacementForceOperator::GetResidualPreparationCount() const noexcept {
        return m_residualPreparationCount;
    }

    std::uint64_t PreparedGravityDisplacementForceOperator::GetResidualApplicationCount() const noexcept {
        return m_residualApplicationCount;
    }

    const PreparedGravityDisplacementForceColumnStatistics &
    PreparedGravityDisplacementForceOperator::GetDensityJacobianStatistics() const noexcept {
        return m_densityJacobianStatistics;
    }

    const PreparedGravityDisplacementForceColumnStatistics &
    PreparedGravityDisplacementForceOperator::GetGravityGradientJacobianStatistics() const noexcept {
        return m_gravityGradientJacobianStatistics;
    }

    const PreparedGravityDisplacementForceColumnStatistics &
    PreparedGravityDisplacementForceOperator::GetDisplacementJacobianStatistics() const noexcept {
        return m_displacementJacobianStatistics;
    }

    const PreparedGravityDisplacementForceCompleteStatistics &
    PreparedGravityDisplacementForceOperator::GetCompleteJacobianStatistics() const noexcept {
        return m_completeJacobianStatistics;
    }

    const fem::FEM &PreparedGravityDisplacementForceOperator::GetFEM() const noexcept {
        return m_fem;
    }

    const context::gravity_field::GravityFieldLinearizationContext &
    PreparedGravityDisplacementForceOperator::GetGravityContext() const noexcept {
        return m_gravityContext;
    }

    void PreparedGravityDisplacementForceOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            IsPrepared(), "PreparedGravityDisplacementForceOperator must be prepared for "
                          "the current shared gravity-context revisions before residual or "
                          "Jacobian application."
        );
    }

    PreparedGravityDisplacementForceJacobianOperator::PreparedGravityDisplacementForceJacobianOperator(
        const GravityDisplacementForceLayout &layout,
        const PreparedGravityDisplacementForceOperator &preparedOperator
    )
        : mfem::Operator(
              layout.residual_offsets().Last(),
              layout.value_offsets().Last()
          ),
          m_layout(layout),
          m_preparedOperator(preparedOperator) {
        using Form                  = utils::blocks::barotropic_equilibrium_form;

        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);

        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);

        constexpr auto gravityGradientValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);

        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);

        MFEM_VERIFY(
            m_layout.size(densityValue) == m_preparedOperator.GetGravityContext().GetDensityMap().reduced_size() &&
                m_layout.size(displacementValue) ==
                    m_preparedOperator.GetGravityContext().GetDisplacementMap().reduced_size() &&
                m_layout.size(gravityGradientValue) ==
                    m_preparedOperator.GetGravityContext().GetGravityGradientMap().reduced_size() &&
                m_layout.size(displacementResidual) ==
                    m_preparedOperator.GetGravityContext().GetDisplacementMap().reduced_size(),
            "Prepared gravity-displacement-force MFEM adapter received "
            "incompatible coupled block sizes."
        );
    }

    void PreparedGravityDisplacementForceJacobianOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_preparedOperator.IsPrepared(), "Prepared gravity-displacement-force MFEM adapter requires a "
                                             "prepared operator."
        );

        MFEM_VERIFY(
            direction.Size() == Width(), "Prepared gravity-displacement-force MFEM adapter received a "
                                         "direction with the wrong size."
        );

        using Form                  = utils::blocks::barotropic_equilibrium_form;

        constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);

        constexpr auto displacementValue =
            utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);

        constexpr auto gravityGradientValue =
            utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);

        constexpr auto displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);

        const mfem::Vector densityVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(densityValue), m_layout.size(densityValue)
        );

        const mfem::Vector displacementVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(displacementValue),
            m_layout.size(displacementValue)
        );

        const mfem::Vector gravityGradientVariation(
            const_cast<mfem::real_t *>(direction.GetData()) + m_layout.offset(gravityGradientValue),
            m_layout.size(gravityGradientValue)
        );

        mfem::Vector displacementAction;

        m_preparedOperator.ApplyCompleteJacobianAction(
            densityVariation, displacementVariation, gravityGradientVariation, displacementAction
        );

        MFEM_VERIFY(
            displacementAction.Size() == m_layout.size(displacementResidual),
            "Prepared gravity-displacement-force MFEM adapter produced a "
            "displacement action with the wrong size."
        );

        action.SetSize(Height());
        action                   = 0.0;

        const int residualOffset = m_layout.offset(displacementResidual);

        for (int entry = 0; entry < displacementAction.Size(); ++entry) {
            action(residualOffset + entry) = displacementAction(entry);
        }
    }

    const GravityDisplacementForceLayout &PreparedGravityDisplacementForceJacobianOperator::GetLayout() const noexcept {
        return m_layout;
    }
} // namespace mean_field::operators