module;

#include <cmath>

#include <mfem.hpp>

module mean_field;

import :operators.context.pressure_force;

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

    template <typename Dependency>
    void validate_dependency_transition(
        const Dependency &prepared,
        const Dependency &requested,
        const char *message
    ) {
        MFEM_VERIFY(requested.CanFollow(prepared), message);

        MFEM_VERIFY(
            prepared.identity == requested.identity || prepared.revision != requested.revision,
            "A new pressure-force dependency identity must also carry "
            "a visibly different revision."
        );
    }
} // namespace

namespace mean_field::operators::context::pressure_force {
    PressureForceLinearizationContext::PressureForceLinearizationContext(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const field::FieldDofMap &enthalpyMap,
        const field::FieldDofMap &displacementMap
    )
        : m_enthalpySize(enthalpyMap.reduced_size()),
          m_displacementSize(displacementMap.reduced_size()) {
        MFEM_VERIFY(f.mesh != nullptr, "PressureForceLinearizationContext requires a mesh.");

        MFEM_VERIFY(
            f.enthalpyFes != nullptr, "PressureForceLinearizationContext requires the enthalpy "
                                      "finite-element space."
        );

        MFEM_VERIFY(
            f.displacementFes != nullptr, "PressureForceLinearizationContext requires the displacement "
                                          "finite-element space."
        );

        MFEM_VERIFY(
            domainMapper.GetDimension() == f.mesh->Dimension(), "The pressure-force context's stateless domain-mapper "
                                                                "dimension does not match the mesh dimension."
        );

        MFEM_VERIFY(
            enthalpyMap.full_size() == f.enthalpyFes->GetTrueVSize(),
            "The pressure-force enthalpy FieldDofMap does not match the "
            "enthalpy finite-element space."
        );

        MFEM_VERIFY(
            displacementMap.full_size() == f.displacementFes->GetTrueVSize(),
            "The pressure-force displacement FieldDofMap does not match "
            "the displacement finite-element space."
        );
    }

    PressureForcePreparationReport PressureForceLinearizationContext::Prepare(
        const PressureForceStateView &state,
        const PressureForceDependencies &dependencies
    ) {
        MFEM_VERIFY(
            state.enthalpy.Size() == m_enthalpySize, "PressureForceLinearizationContext received a supported "
                                                     "enthalpy vector with the wrong size."
        );

        MFEM_VERIFY(
            state.displacement.Size() == m_displacementSize, "PressureForceLinearizationContext received a supported "
                                                             "displacement vector with the wrong size."
        );

        validate_finite_vector(
            state.enthalpy, "PressureForceLinearizationContext received a non-finite "
                            "enthalpy value."
        );

        validate_finite_vector(
            state.displacement, "PressureForceLinearizationContext received a non-finite "
                                "displacement value."
        );

        if (m_isPrepared) {
            validate_dependency_transition(
                m_dependencies.discretization, dependencies.discretization,
                "PressureForceLinearizationContext received an older "
                "discretization revision for the same identity."
            );

            validate_dependency_transition(
                m_dependencies.enthalpy, dependencies.enthalpy,
                "PressureForceLinearizationContext received an older "
                "enthalpy revision for the same identity."
            );

            validate_dependency_transition(
                m_dependencies.displacement, dependencies.displacement,
                "PressureForceLinearizationContext received an older "
                "displacement revision for the same identity."
            );
        }

        const bool discretizationChanged =
            !m_isPrepared || dependencies.discretization != m_dependencies.discretization;

        const bool enthalpyChanged     = !m_isPrepared || dependencies.enthalpy != m_dependencies.enthalpy;

        const bool displacementChanged = !m_isPrepared || dependencies.displacement != m_dependencies.displacement;

        /*
         * Static data depend only on discretization.
         *
         * Geometry data depend on discretization and displacement.
         *
         * Material data depend on both geometry and enthalpy because
         * pressure and its enthalpy derivative are evaluated on the frozen
         * mapped state.
         */
        const bool geometryPreparationRequired = discretizationChanged || displacementChanged;

        const bool materialPreparationRequired = geometryPreparationRequired || enthalpyChanged;

        PressureForcePreparationReport report;

        report.preparedStaticDependencies = discretizationChanged;

        report.preparedGeometryState      = geometryPreparationRequired;

        report.preparedMaterialState      = materialPreparationRequired;

        /*
         * A discretization change invalidates every frozen field because
         * their coordinate interpretation may have changed.
         */
        if (discretizationChanged || 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.preparedMaterialState) {
            ++m_statistics.materialPreparations;
        }

        m_dependencies = dependencies;

        m_isPrepared   = true;

        return report;
    }

    const PressureForcePreparationStatistics &
    PressureForceLinearizationContext::GetPreparationStatistics() const noexcept {
        return m_statistics;
    }

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

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

    const PressureForceDependencies &PressureForceLinearizationContext::GetDependencies() const {
        VerifyPrepared();

        return m_dependencies;
    }

    const mfem::Vector &PressureForceLinearizationContext::GetBaseEnthalpy() const {
        VerifyPrepared();

        return m_baseEnthalpy;
    }

    const mfem::Vector &PressureForceLinearizationContext::GetDisplacement() const {
        VerifyPrepared();

        return m_displacement;
    }

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