module;

#include <cstdint>

#include <mfem.hpp>

module mean_field;
import :operators.context.barotropic_closure_linearization;

namespace mean_field::operators::context::barotropic {
    BarotropicClosureLinearizationContext::
        BarotropicClosureLinearizationContext(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const physics::PolytropicBarotrope &barotrope
        )
        : m_f(f),
          m_operator(
              f,
              domainMapper,
              barotrope
          ) {
        MFEM_VERIFY(
            m_f.densityFes != nullptr,
            "The closure linearization context requires the "
            "density finite-element space."
        );

        MFEM_VERIFY(
            m_f.enthalpyFes != nullptr,
            "The closure linearization context requires the "
            "enthalpy finite-element space."
        );

        MFEM_VERIFY(
            m_f.displacementFes != nullptr,
            "The closure linearization context requires the "
            "displacement finite-element space."
        );
    }

    void BarotropicClosureLinearizationContext::Prepare(
        const mfem::Vector &baseDensityTrue,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &displacementTrue,
        const BarotropicClosureRevisions &revisions
    ) {
        MFEM_VERIFY(
            baseDensityTrue.Size() == m_f.densityFes->GetTrueVSize(),
            "The closure base-density vector has the wrong size."
        );

        MFEM_VERIFY(
            baseEnthalpyTrue.Size() == m_f.enthalpyFes->GetTrueVSize(),
            "The closure base-enthalpy vector has the wrong size."
        );

        MFEM_VERIFY(
            displacementTrue.Size() == m_f.displacementFes->GetTrueVSize(),
            "The closure displacement vector has the wrong size."
        );

        if (m_isPrepared && revisions == m_revisions) {
            return;
        }

        m_operator.Prepare(baseDensityTrue, baseEnthalpyTrue, displacementTrue);

        m_baseDensityTrue  = baseDensityTrue;
        m_baseEnthalpyTrue = baseEnthalpyTrue;
        m_displacementTrue = displacementTrue;

        m_revisions        = revisions;
        m_isPrepared       = true;
        ++m_preparationCount;
    }

    bool BarotropicClosureLinearizationContext::IsPrepared() const noexcept {
        return m_isPrepared;
    }

    bool BarotropicClosureLinearizationContext::MatchesRevisions(
        const BarotropicClosureRevisions &revisions
    ) const noexcept {
        return m_isPrepared && revisions == m_revisions;
    }

    std::uint64_t BarotropicClosureLinearizationContext::
        GetPreparationCount() const noexcept {
        return m_preparationCount;
    }

    const BarotropicClosureRevisions &
    BarotropicClosureLinearizationContext::GetRevisions() const {
        VerifyPrepared();
        return m_revisions;
    }

    const mfem::Vector &
    BarotropicClosureLinearizationContext::GetBaseDensityTrue() const {
        VerifyPrepared();
        return m_baseDensityTrue;
    }

    const mfem::Vector &
    BarotropicClosureLinearizationContext::GetBaseEnthalpyTrue() const {
        VerifyPrepared();
        return m_baseEnthalpyTrue;
    }

    const mfem::Vector &
    BarotropicClosureLinearizationContext::GetDisplacementTrue() const {
        VerifyPrepared();
        return m_displacementTrue;
    }

    const PreparedBarotropicClosureOperator &
    BarotropicClosureLinearizationContext::GetOperator() const noexcept {
        return m_operator;
    }

    void BarotropicClosureLinearizationContext::BuildResidual(
        mfem::Vector &residual
    ) const {
        VerifyPrepared();
        m_operator.BuildResidual(residual);
    }

    void BarotropicClosureLinearizationContext::VerifyPrepared() const {
        MFEM_VERIFY(
            m_isPrepared, "The barotropic-closure linearization context "
                          "has not been prepared."
        );
    }
} // namespace mean_field::operators::context::barotropic