module;

#include <mfem.hpp>

module mean_field;

import :operators.kernels.rotational_displacement_force;
import :operators.prepared_rotational_displacement_force;

namespace mean_field::operators {
    PreparedRotationalDisplacementForceOperator::PreparedRotationalDisplacementForceOperator(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper
    )
        : m_fem(f),
          m_domainMapper(domainMapper),
          m_context(
              f,
              domainMapper
          ) {
        MFEM_VERIFY(m_fem.mesh != nullptr, "PreparedRotationalDisplacementForceOperator requires a mesh.");

        MFEM_VERIFY(
            m_fem.mesh->Dimension() == 3, "PreparedRotationalDisplacementForceOperator requires a "
                                          "three-dimensional mesh."
        );

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

        MFEM_VERIFY(
            m_fem.compactificationFes != nullptr && m_fem.compactificationCoordinate != nullptr,
            "PreparedRotationalDisplacementForceOperator requires the "
            "compactification coordinate."
        );

        MFEM_VERIFY(
            m_fem.quadratureFactory != nullptr, "PreparedRotationalDisplacementForceOperator requires the "
                                                "quadrature-rule factory."
        );

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

    PreparedRotationalDisplacementForceReport PreparedRotationalDisplacementForceOperator::Prepare(
        const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
        const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
        const physics::RigidRotation &rotation
    ) {
        const bool rotationChanged =
            !m_context.IsPrepared() || dependencies.rotation != m_context.GetDependencies().rotation;

        PreparedRotationalDisplacementForceReport report;
        report.contextReport = m_context.Prepare(state, dependencies);

        if (!report.contextReport.DidAnyWork()) {
            return report;
        }

        m_isPrepared = false;

        if (rotationChanged) {
            m_rotation             = rotation;
            report.updatedRotation = true;
        }

        MFEM_VERIFY(
            m_rotation.has_value(), "PreparedRotationalDisplacementForceOperator has no frozen "
                                    "rotation state."
        );

        if (report.contextReport.preparedBaseState) {
            kernels::apply_rotational_displacement_force_residual(
                m_fem, m_domainMapper, *m_rotation, m_context.GetBaseDensityTrue(), m_context.GetDisplacementTrue(),
                m_actionTrue
            );
            m_cachedResidual.SetSize(m_context.GetDisplacementMap().reduced_size());
            m_context.GetDisplacementMap().gather(m_actionTrue, m_cachedResidual);

            ++m_residualPreparationCount;
            report.preparedResidual = true;
        }

        MFEM_VERIFY(
            m_cachedResidual.Size() == m_context.GetDisplacementMap().reduced_size(),
            "The prepared rotational-displacement-force residual has the "
            "wrong size."
        );

        m_preparedDependencies = dependencies;
        m_isPrepared           = true;

        return report;
    }

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

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

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

        kernels::apply_rotational_displacement_force_density_action(
            m_fem, m_domainMapper, *m_rotation, m_densityVariationTrue, m_context.GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_densityJacobianStatistics.applications;
    }

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

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

        kernels::apply_rotational_displacement_force_displacement_action(
            m_fem, m_domainMapper, *m_rotation, m_context.GetBaseDensityTrue(), m_displacementVariationTrue,
            m_context.GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

        ++m_displacementJacobianStatistics.applications;
    }

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

        m_densityVariationTrue.SetSize(m_context.GetDensityMap().full_size());
        m_displacementVariationTrue.SetSize(m_context.GetDisplacementMap().full_size());
        m_context.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
        m_context.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);

        kernels::apply_rotational_displacement_force_complete_action(
            m_fem, m_domainMapper, *m_rotation, m_context.GetBaseDensityTrue(), m_densityVariationTrue,
            m_displacementVariationTrue, m_context.GetDisplacementTrue(), m_actionTrue
        );
        action.SetSize(m_context.GetDisplacementMap().reduced_size());
        m_context.GetDisplacementMap().gather(m_actionTrue, action);

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

    bool PreparedRotationalDisplacementForceOperator::IsPrepared() const noexcept {
        return m_isPrepared && m_rotation.has_value() && m_context.MatchesDependencies(m_preparedDependencies);
    }

    const context::rotational_displacement_force::RotationalDisplacementForcePreparationStatistics &
    PreparedRotationalDisplacementForceOperator::GetContextPreparationStatistics() const noexcept {
        return m_context.GetPreparationStatistics();
    }

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

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

    const PreparedRotationalDisplacementForceColumnStatistics &
    PreparedRotationalDisplacementForceOperator::GetDensityJacobianStatistics() const noexcept {
        return m_densityJacobianStatistics;
    }

    const PreparedRotationalDisplacementForceColumnStatistics &
    PreparedRotationalDisplacementForceOperator::GetDisplacementJacobianStatistics() const noexcept {
        return m_displacementJacobianStatistics;
    }

    const PreparedRotationalDisplacementForceCompleteStatistics &
    PreparedRotationalDisplacementForceOperator::GetCompleteJacobianStatistics() const noexcept {
        return m_completeJacobianStatistics;
    }

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

    const context::rotational_displacement_force::RotationalDisplacementForceLinearizationContext &
    PreparedRotationalDisplacementForceOperator::GetContext() const noexcept {
        return m_context;
    }

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

    PreparedRotationalDisplacementForceJacobianOperator::PreparedRotationalDisplacementForceJacobianOperator(
        const RotationalDisplacementForceLayout &layout,
        const PreparedRotationalDisplacementForceOperator &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 displacementResidual =
            utils::blocks::get_residual_block<Form>(utils::blocks::displacement_field.geometry_term);

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

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

        MFEM_VERIFY(
            direction.Size() == Width(), "Prepared rotational-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 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)
        );

        mfem::Vector displacementAction;

        m_preparedOperator.ApplyCompleteJacobianAction(densityVariation, displacementVariation, displacementAction);

        MFEM_VERIFY(
            displacementAction.Size() == m_layout.size(displacementResidual),
            "Prepared rotational-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 RotationalDisplacementForceLayout &
    PreparedRotationalDisplacementForceJacobianOperator::GetLayout() const noexcept {
        return m_layout;
    }
} // namespace mean_field::operators