module;

#include <cmath>

#include <mfem.hpp>

module mean_field;

import :operators.context.hydrostatic_equilibrium;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

    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);
        }
    }

    void validate_state(
        const mean_field::field::FieldDofMap &enthalpyMap,
        const mean_field::field::FieldDofMap &gravityPotentialMap,
        const mean_field::field::FieldDofMap &displacementMap,
        const mean_field::operators::context::hydrostatic::HydrostaticEquilibriumStateView &state
    ) {
        MFEM_VERIFY(
            state.enthalpy.Size() == enthalpyMap.reduced_size(), "HydrostaticEquilibriumContext received a supported "
                                                                 "enthalpy vector with the wrong size."
        );

        MFEM_VERIFY(
            state.gravityPotential.Size() == gravityPotentialMap.reduced_size(),
            "HydrostaticEquilibriumContext received a supported gravity-potential vector with the wrong size."
        );

        MFEM_VERIFY(
            state.displacement.Size() == displacementMap.reduced_size(),
            "HydrostaticEquilibriumContext received a supported displacement vector with the wrong size."
        );

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

        validate_finite_vector(
            state.gravityPotential, "HydrostaticEquilibriumContext received a "
                                    "non-finite gravity-potential value."
        );

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

        MFEM_VERIFY(
            std::isfinite(state.bernoulliConstant), "HydrostaticEquilibriumContext received a "
                                                    "non-finite Bernoulli constant."
        );
    }

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

namespace mean_field::operators::context::hydrostatic {
    HydrostaticEquilibriumContext::HydrostaticEquilibriumContext(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper
    )
        : m_f(f),
          m_domainMapper(domainMapper),
          m_enthalpyMap(
              field::make_field_dof_map<field::Enthalpy, DomainSchema>(*f.enthalpyFes)
          ),
          m_gravityPotentialMap(
              field::make_field_dof_map<field::Gravity, DomainSchema>(*f.gravityPotentialFes)
          ),
          m_displacementMap(
              field::make_field_dof_map<field::Displacement, DomainSchema>(*f.displacementFes)
          ) {
        MFEM_VERIFY(m_f.mesh != nullptr, "HydrostaticEquilibriumContext requires a mesh.");

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

        MFEM_VERIFY(
            m_f.gravityPotentialFes != nullptr, "HydrostaticEquilibriumContext requires the "
                                                "gravity-potential finite-element space."
        );

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

        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_f.mesh->Dimension(), "The hydrostatic context's stateless "
                                                                    "domain-mapper dimension does not match the mesh "
                                                                    "dimension."
        );

        m_baseEnthalpyTrue.SetSize(m_enthalpyMap.full_size());
        m_baseGravityPotentialTrue.SetSize(m_gravityPotentialMap.full_size());
        m_displacementTrue.SetSize(m_displacementMap.full_size());
    }

    HydrostaticPreparationReport HydrostaticEquilibriumContext::Prepare(
        const HydrostaticEquilibriumStateView &state,
        const HydrostaticEquilibriumDependencies &dependencies
    ) {
        validate_state(m_enthalpyMap, m_gravityPotentialMap, m_displacementMap, state);

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

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

            validate_dependency_transition(
                m_dependencies.gravityPotential, dependencies.gravityPotential,
                "HydrostaticEquilibriumContext received an older "
                "gravity-potential revision for the same identity."
            );

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

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

            validate_dependency_transition(
                m_dependencies.bernoulliConstant, dependencies.bernoulliConstant,
                "HydrostaticEquilibriumContext received an older "
                "Bernoulli-constant revision for the same identity."
            );
        }

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

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

        const bool gravityPotentialChanged =
            !m_isPrepared || dependencies.gravityPotential != m_dependencies.gravityPotential;

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

        const bool rotationChanged     = !m_isPrepared || dependencies.rotation != m_dependencies.rotation;

        const bool bernoulliConstantChanged =
            !m_isPrepared || dependencies.bernoulliConstant != m_dependencies.bernoulliConstant;

        const bool geometryPreparationRequired = staticChanged || displacementChanged;

        const bool rotationPreparationRequired = geometryPreparationRequired || rotationChanged;

        const bool baseStatePreparationRequired =
            rotationPreparationRequired || enthalpyChanged || gravityPotentialChanged || bernoulliConstantChanged;

        HydrostaticPreparationReport report;

        report.preparedStaticDependencies   = staticChanged;
        report.preparedGeometryState        = geometryPreparationRequired;
        report.preparedRotationDependencies = rotationPreparationRequired;
        report.preparedBaseState            = baseStatePreparationRequired;

        if (staticChanged || enthalpyChanged) {
            m_enthalpyMap.scatter(state.enthalpy, m_baseEnthalpyTrue);
            report.updatedEnthalpy = true;
        }

        if (staticChanged || gravityPotentialChanged) {
            m_gravityPotentialMap.scatter(state.gravityPotential, m_baseGravityPotentialTrue);

            report.updatedGravityPotential = true;
        }

        if (geometryPreparationRequired) {
            m_displacementMap.scatter(state.displacement, m_displacementTrue);
            report.updatedDisplacement = true;
        }

        if (staticChanged || bernoulliConstantChanged) {
            m_bernoulliConstant             = state.bernoulliConstant;
            report.updatedBernoulliConstant = true;
        }

        if (report.preparedStaticDependencies) {
            ++m_statistics.staticPreparations;
        }

        if (report.preparedGeometryState) {
            ++m_statistics.geometryPreparations;
        }

        if (report.preparedRotationDependencies) {
            ++m_statistics.rotationPreparations;
        }

        if (report.preparedBaseState) {
            ++m_statistics.baseStatePreparations;
        }

        m_dependencies = dependencies;
        m_isPrepared   = true;

        return report;
    }

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

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

    const HydrostaticEquilibriumDependencies &HydrostaticEquilibriumContext::GetDependencies() const {
        VerifyPrepared();
        return m_dependencies;
    }

    const HydrostaticPreparationStatistics &HydrostaticEquilibriumContext::GetPreparationStatistics() const noexcept {
        return m_statistics;
    }

    const field::FieldDofMap &HydrostaticEquilibriumContext::GetEnthalpyMap() const noexcept {
        return m_enthalpyMap;
    }

    const field::FieldDofMap &HydrostaticEquilibriumContext::GetGravityPotentialMap() const noexcept {
        return m_gravityPotentialMap;
    }

    const field::FieldDofMap &HydrostaticEquilibriumContext::GetDisplacementMap() const noexcept {
        return m_displacementMap;
    }

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

    const mfem::Vector &HydrostaticEquilibriumContext::GetBaseGravityPotentialTrue() const {
        VerifyPrepared();
        return m_baseGravityPotentialTrue;
    }

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

    double HydrostaticEquilibriumContext::GetBernoulliConstant() const {
        VerifyPrepared();
        return m_bernoulliConstant;
    }

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