module;

#include <cmath>

#include <mfem.hpp>

module mean_field;

import :operators.context.barotropic_closure_linearization;

namespace {
    void validate_finite_vector(
        const mfem::Vector &vector,
        const char *message
    ) {
        for (int i = 0; i < vector.Size(); ++i) {
            MFEM_VERIFY(std::isfinite(vector(i)), message);
        }
    }

    template <typename Stamp>
    void validate_dependency_transition(
        const Stamp &prepared,
        const Stamp &requested,
        const char *message
    ) {
        MFEM_VERIFY(requested.CanFollow(prepared), message);
        MFEM_VERIFY(
            prepared.identity == requested.identity || prepared.revision != requested.revision,
            "A new barotropic-closure dependency identity must also carry a visibly different revision."
        );
    }
} // namespace

namespace mean_field::operators::context::barotropic {
    BarotropicClosureLinearizationContext::BarotropicClosureLinearizationContext(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const field::FieldDofMap &densityMap,
        const field::FieldDofMap &enthalpyMap,
        const field::FieldDofMap &displacementMap
    )
        : m_f(f),
          m_domainMapper(domainMapper),
          m_densitySize(densityMap.reduced_size()),
          m_enthalpySize(enthalpyMap.reduced_size()),
          m_displacementSize(displacementMap.reduced_size()) {
        MFEM_VERIFY(m_f.mesh != nullptr, "BarotropicClosureLinearizationContext requires a mesh.");
        MFEM_VERIFY(m_f.densityFes != nullptr, "BarotropicClosureLinearizationContext requires the density FE space.");
        MFEM_VERIFY(
            m_f.enthalpyFes != nullptr, "BarotropicClosureLinearizationContext requires the enthalpy FE space."
        );
        MFEM_VERIFY(
            m_f.displacementFes != nullptr, "BarotropicClosureLinearizationContext requires the displacement FE space."
        );
        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_f.mesh->Dimension(),
            "The barotropic-closure context domain-mapper dimension does not match the mesh dimension."
        );
        MFEM_VERIFY(
            densityMap.full_size() == m_f.densityFes->GetTrueVSize(),
            "The density FieldDofMap does not match the density FE space."
        );
        MFEM_VERIFY(
            enthalpyMap.full_size() == m_f.enthalpyFes->GetTrueVSize(),
            "The enthalpy FieldDofMap does not match the enthalpy FE space."
        );
        MFEM_VERIFY(
            displacementMap.full_size() == m_f.displacementFes->GetTrueVSize(),
            "The displacement FieldDofMap does not match the displacement FE space."
        );
    }

    BarotropicClosurePreparationReport BarotropicClosureLinearizationContext::Prepare(
        const BarotropicClosureStateView &state,
        const BarotropicClosureDependencies &dependencies
    ) {
        MFEM_VERIFY(state.density.Size() == m_densitySize, "The supported closure density vector has the wrong size.");
        MFEM_VERIFY(
            state.enthalpy.Size() == m_enthalpySize, "The supported closure enthalpy vector has the wrong size."
        );
        MFEM_VERIFY(
            state.displacement.Size() == m_displacementSize,
            "The supported closure displacement vector has the wrong size."
        );

        validate_finite_vector(state.density, "The closure density state contains a non-finite value.");
        validate_finite_vector(state.enthalpy, "The closure enthalpy state contains a non-finite value.");
        validate_finite_vector(state.displacement, "The closure displacement state contains a non-finite value.");

        if (m_isPrepared) {
            validate_dependency_transition(
                m_dependencies.discretization, dependencies.discretization,
                "BarotropicClosureLinearizationContext received an older discretization revision for the same identity."
            );
            validate_dependency_transition(
                m_dependencies.density, dependencies.density,
                "BarotropicClosureLinearizationContext received an older density revision for the same identity."
            );
            validate_dependency_transition(
                m_dependencies.enthalpy, dependencies.enthalpy,
                "BarotropicClosureLinearizationContext received an older enthalpy revision for the same identity."
            );
            validate_dependency_transition(
                m_dependencies.displacement, dependencies.displacement,
                "BarotropicClosureLinearizationContext received an older displacement revision for the same identity."
            );
        }

        const bool staticChanged       = !m_isPrepared || dependencies.discretization != m_dependencies.discretization;
        const bool densityChanged      = !m_isPrepared || dependencies.density != m_dependencies.density;
        const bool enthalpyChanged     = !m_isPrepared || dependencies.enthalpy != m_dependencies.enthalpy;
        const bool displacementChanged = !m_isPrepared || dependencies.displacement != m_dependencies.displacement;

        const bool geometryPreparationRequired = staticChanged || displacementChanged;
        const bool baseStatePreparationRequired =
            staticChanged || geometryPreparationRequired || densityChanged || enthalpyChanged;

        BarotropicClosurePreparationReport report;
        report.preparedStaticDependencies = staticChanged;
        report.preparedGeometryState      = geometryPreparationRequired;
        report.preparedBaseState          = baseStatePreparationRequired;

        if (staticChanged || densityChanged) {
            m_baseDensity         = state.density;
            report.updatedDensity = true;
        }
        if (staticChanged || enthalpyChanged) {
            m_baseEnthalpy         = state.enthalpy;
            report.updatedEnthalpy = true;
        }
        if (geometryPreparationRequired) {
            m_displacement             = state.displacement;
            report.updatedDisplacement = true;
        }

        if (report.preparedStaticDependencies) {
            ++m_statistics.staticPreparations;
        }
        if (report.preparedGeometryState) {
            ++m_statistics.geometryPreparations;
        }
        if (report.preparedBaseState) {
            ++m_statistics.baseStatePreparations;
        }

        m_dependencies = dependencies;
        m_isPrepared   = true;
        return report;
    }

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

    bool BarotropicClosureLinearizationContext::MatchesDependencies(
        const BarotropicClosureDependencies &dependencies
    ) const noexcept {
        return m_isPrepared && dependencies == m_dependencies;
    }

    const BarotropicClosureDependencies &BarotropicClosureLinearizationContext::GetDependencies() const {
        VerifyPrepared();
        return m_dependencies;
    }

    const BarotropicClosurePreparationStatistics &
    BarotropicClosureLinearizationContext::GetPreparationStatistics() const noexcept {
        return m_statistics;
    }

    const mfem::Vector &BarotropicClosureLinearizationContext::GetBaseDensity() const {
        VerifyPrepared();
        return m_baseDensity;
    }

    const mfem::Vector &BarotropicClosureLinearizationContext::GetBaseEnthalpy() const {
        VerifyPrepared();
        return m_baseEnthalpy;
    }

    const mfem::Vector &BarotropicClosureLinearizationContext::GetDisplacement() const {
        VerifyPrepared();
        return m_displacement;
    }

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