module;

#include <cmath>
#include <utility>

#include <mfem.hpp>

export module mean_field:operators.prepared_centering_constraint;

export import :field.mfem;

export namespace mean_field::operators {
    struct PreparedCenteringConstraintReport final {
        bool cachedCenterDisplacement{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return cachedCenterDisplacement;
        }
    };

    /*
     * Strong translational gauge: the material point at the computational
     * origin has zero displacement. The three corresponding displacement
     * residual rows replace redundant force-balance rows.
     */
    class PreparedCenteringConstraint final {
    public:
        explicit PreparedCenteringConstraint(field::FieldPointDofMap centerRows)
            : m_centerRows(std::move(centerRows)),
              m_centerDisplacement(m_centerRows.size()) {
        }

        [[nodiscard]] PreparedCenteringConstraintReport Prepare(
            const mfem::Vector &displacement,
            const bool displacementChanged
        ) {
            MFEM_VERIFY(
                displacement.Size() == m_centerRows.field_size(),
                "The centering constraint received a displacement vector with the wrong size."
            );

            PreparedCenteringConstraintReport report;
            if (!m_isPrepared || displacementChanged) {
                for (int centerIndex = 0; centerIndex < m_centerRows.size(); ++centerIndex) {
                    const double value = displacement(m_centerRows.reduced_dofs()[centerIndex]);
                    MFEM_VERIFY(
                        std::isfinite(value), "The centering constraint received a non-finite center displacement."
                    );
                    m_centerDisplacement(centerIndex) = value;
                }
                report.cachedCenterDisplacement = true;
            }

            m_isPrepared = true;
            return report;
        }

        void ApplyResidualRows(mfem::Vector &displacementResidual) const {
            VerifyPrepared();
            MFEM_VERIFY(
                displacementResidual.Size() == m_centerRows.field_size(),
                "The centering constraint received a displacement residual with the wrong size."
            );

            for (int centerIndex = 0; centerIndex < m_centerRows.size(); ++centerIndex) {
                displacementResidual(m_centerRows.reduced_dofs()[centerIndex]) = m_centerDisplacement(centerIndex);
            }
        }

        void ApplyJacobianRows(
            const mfem::Vector &displacementVariation,
            mfem::Vector &displacementAction
        ) const {
            VerifyPrepared();
            MFEM_VERIFY(
                displacementVariation.Size() == m_centerRows.field_size() &&
                    displacementAction.Size() == m_centerRows.field_size(),
                "The centering constraint received a Jacobian vector with the wrong size."
            );

            for (const int centerRow : m_centerRows.reduced_dofs()) {
                displacementAction(centerRow) = displacementVariation(centerRow);
            }
        }

        [[nodiscard]] bool IsPrepared() const noexcept {
            return m_isPrepared;
        }

        [[nodiscard]] const field::FieldPointDofMap &GetCenterRows() const noexcept {
            return m_centerRows;
        }

    private:
        void VerifyPrepared() const {
            MFEM_VERIFY(m_isPrepared, "The centering constraint must be prepared before row application.");
        }

        field::FieldPointDofMap m_centerRows;
        mfem::Vector m_centerDisplacement;
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators