module;

#include <array>
#include <cmath>
#include <cstdint>
#include <limits>
#include <mfem.hpp>

module mean_field;

import :operators.prepared_barotropic_closure;

namespace {
    int get_density_size(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.densityFes != nullptr,
            "PreparedBarotropicClosureOperator requires the "
            "density finite-element space."
        );

        return f.densityFes->GetTrueVSize();
    }

    int get_enthalpy_size(const mean_field::fem::FEM &f) {
        MFEM_VERIFY(
            f.enthalpyFes != nullptr,
            "PreparedBarotropicClosureOperator requires the "
            "enthalpy finite-element space."
        );

        return f.enthalpyFes->GetTrueVSize();
    }

    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 true_to_local(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &trueVector,
        mfem::Vector &localVector
    ) {
        MFEM_VERIFY(
            trueVector.Size() == finiteElementSpace.GetTrueVSize(),
            "True vector has the wrong size."
        );

        localVector.SetSize(finiteElementSpace.GetVSize());

        const mfem::Operator *prolongation =
            finiteElementSpace.GetProlongationMatrix();

        if (prolongation != nullptr) {
            prolongation->Mult(trueVector, localVector);
        } else {
            localVector = trueVector;
        }
    }

    void local_to_true(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &localVector,
        mfem::Vector &trueVector
    ) {
        MFEM_VERIFY(
            localVector.Size() == finiteElementSpace.GetVSize(),
            "Local vector has the wrong size."
        );

        trueVector.SetSize(finiteElementSpace.GetTrueVSize());
        trueVector = 0.0;

        const mfem::Operator *prolongation =
            finiteElementSpace.GetProlongationMatrix();

        if (prolongation != nullptr) {
            prolongation->MultTranspose(localVector, trueVector);
        } else {
            trueVector = localVector;
        }
    }

    int get_eos_extra_order(
        const mean_field::physics::PolytropicBarotrope &barotrope
    ) {
        const double extraOrder =
            (barotrope.polytropic_index() - 1.0) *
            static_cast<double>(
                mean_field::field::Enthalpy::Scalar::familyOrder
            );

        MFEM_VERIFY(
            std::isfinite(extraOrder) && extraOrder >= 0.0 &&
                extraOrder <=
                    static_cast<double>(std::numeric_limits<int>::max()),
            "The EOS effective polynomial order is invalid."
        );

        return static_cast<int>(std::ceil(extraOrder));
    }

    const mfem::IntegrationRule &get_eos_rule(
        const mean_field::fem::FEM &f,
        const mean_field::physics::PolytropicBarotrope &barotrope,
        const mfem::FiniteElement &densityElement,
        const mfem::FiniteElement &enthalpyElement,
        const mfem::ElementTransformation &transformation
    ) {
        using EnthalpyField =
            mean_field::field::Field<mean_field::field::Enthalpy>;

        MFEM_VERIFY(
            densityElement.GetOrder() ==
                mean_field::field::Density::Scalar::familyOrder,
            "The prepared EOS test element does not match "
            "the registered density field."
        );

        MFEM_VERIFY(
            enthalpyElement.GetOrder() ==
                mean_field::field::Enthalpy::Scalar::familyOrder,
            "The prepared EOS trial element does not match "
            "the registered enthalpy field."
        );

        const mean_field::quadrature::Query query = EnthalpyField::make_query<
            mean_field::field::Enthalpy::Form::EosClosureSource>(
            mean_field::quadrature::QuadratureRole::discretization,
            transformation.OrderW(),
            std::array<int, 1>{get_eos_extra_order(barotrope)},
            mean_field::utils::DOMAINS::STELLAR,
            mean_field::quadrature::MappingKind::general
        );

        const auto resolution =
            f.quadratureFactory->get(query, transformation.GetGeometryType());

        MFEM_VERIFY(
            resolution.integration_rule != nullptr,
            "The quadrature policy did not return a prepared "
            "EOS-closure integration rule."
        );

        return *resolution.integration_rule;
    }
} // namespace

namespace mean_field::operators {
    PreparedBarotropicClosureOperator::PreparedBarotropicClosureOperator(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const physics::PolytropicBarotrope &barotrope
    )
        : mfem::Operator(
              f.densityFes->GetTrueVSize(),
              f.densityFes->GetTrueVSize() + f.enthalpyFes->GetTrueVSize() +
                  f.displacementFes->GetTrueVSize()
          ),
          m_fem(f),
          m_domainMapper(domainMapper),
          m_barotrope(barotrope),
          m_densitySize(f.densityFes->GetTrueVSize()),
          m_enthalpySize(f.enthalpyFes->GetTrueVSize()) {
        MFEM_VERIFY(
            m_fem.densityFes != nullptr,
            "PreparedBarotropicClosureOperator requires "
            "a density finite-element space."
        );

        MFEM_VERIFY(
            m_fem.enthalpyFes != nullptr,
            "PreparedBarotropicClosureOperator requires "
            "an enthalpy finite-element space."
        );

        MFEM_VERIFY(
            m_fem.displacementFes != nullptr,
            "PreparedBarotropicClosureOperator requires "
            "a displacement finite-element space."
        );
    }

    void PreparedBarotropicClosureOperator::Prepare(
        const mfem::Vector &baseDensityTrue,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &displacementTrue
    ) {
        MFEM_VERIFY(
            baseDensityTrue.Size() == m_densitySize,
            "PreparedBarotropicClosureOperator received a "
            "base-density vector with the wrong size."
        );

        MFEM_VERIFY(
            baseEnthalpyTrue.Size() == m_enthalpySize,
            "PreparedBarotropicClosureOperator received a "
            "base-enthalpy vector with the wrong size."
        );

        MFEM_VERIFY(
            displacementTrue.Size() == m_fem.displacementFes->GetTrueVSize(),
            "PreparedBarotropicClosureOperator received a "
            "displacement vector with the wrong size."
        );
        MFEM_VERIFY(
            baseDensityTrue.Size() == m_fem.densityFes->GetTrueVSize(),
            "The base density true vector has the wrong size."
        );

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

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

        validate_finite_vector(
            baseDensityTrue, "PreparedBarotropicClosureOperator received a "
                             "non-finite base-density value."
        );

        validate_finite_vector(
            baseEnthalpyTrue, "PreparedBarotropicClosureOperator received a "
                              "non-finite base-enthalpy value."
        );

        validate_finite_vector(
            displacementTrue, "PreparedBarotropicClosureOperator received a "
                              "non-finite displacement value."
        );

        m_isPrepared = false;
        m_elements.clear();
        m_elements.reserve(m_fem.mesh->GetNE());

        mfem::Vector baseDensityLocal;
        mfem::Vector baseEnthalpyLocal;
        mfem::Vector displacementLocal;

        true_to_local(*m_fem.densityFes, baseDensityTrue, baseDensityLocal);

        true_to_local(*m_fem.enthalpyFes, baseEnthalpyTrue, baseEnthalpyLocal);

        true_to_local(
            *m_fem.displacementFes, displacementTrue, displacementLocal
        );

        mapping::DomainMapperStateless::Workspace workspace(
            m_fem.mesh->Dimension()
        );

        mfem::Array<int> displacementDofs;
        mfem::Array<int> compactificationDofs;

        mfem::Vector elementBaseDensity;
        mfem::Vector elementBaseEnthalpy;
        mfem::Vector elementDisplacement;
        mfem::Vector elementCompactification;

        mfem::Vector densityShape;
        mfem::Vector enthalpyShape;

        const int vacuumAttribute = m_domainMapper.GetVacuumElementAttribute();

        for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
            mfem::ElementTransformation *transformation =
                m_fem.mesh->GetElementTransformation(elementId);

            MFEM_VERIFY(
                transformation != nullptr,
                "PreparedBarotropicClosureOperator received "
                "a null element transformation."
            );

            if (transformation->Attribute == vacuumAttribute) {
                continue;
            }

            m_elements.emplace_back();
            ElementPAData &data = m_elements.back();

            data.densityDofTransformation =
                m_fem.densityFes->GetElementDofs(elementId, data.densityDofs);

            data.enthalpyDofTransformation =
                m_fem.enthalpyFes->GetElementDofs(elementId, data.enthalpyDofs);

            mfem::DofTransformation *displacementDofTransformation =
                m_fem.displacementFes->GetElementVDofs(
                    elementId, displacementDofs
                );

            mfem::DofTransformation *compactificationDofTransformation =
                m_fem.compactificationFes->GetElementDofs(
                    elementId, compactificationDofs
                );

            baseDensityLocal.GetSubVector(data.densityDofs, elementBaseDensity);

            baseEnthalpyLocal.GetSubVector(
                data.enthalpyDofs, elementBaseEnthalpy
            );

            displacementLocal.GetSubVector(
                displacementDofs, elementDisplacement
            );

            m_fem.compactificationCoordinate->GetSubVector(
                compactificationDofs, elementCompactification
            );

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(
                    elementBaseDensity
                );
            }

            if (data.enthalpyDofTransformation != nullptr) {
                data.enthalpyDofTransformation->InvTransformPrimal(
                    elementBaseEnthalpy
                );
            }

            if (displacementDofTransformation != nullptr) {
                displacementDofTransformation->InvTransformPrimal(
                    elementDisplacement
                );
            }

            if (compactificationDofTransformation != nullptr) {
                compactificationDofTransformation->InvTransformPrimal(
                    elementCompactification
                );
            }

            const mfem::FiniteElement &densityElement =
                *m_fem.densityFes->GetFE(elementId);

            const mfem::FiniteElement &enthalpyElement =
                *m_fem.enthalpyFes->GetFE(elementId);

            const mfem::FiniteElement &displacementElement =
                *m_fem.displacementFes->GetFE(elementId);

            const mfem::FiniteElement &compactificationElement =
                *m_fem.compactificationFes->GetFE(elementId);

            const mapping::ElementDisplacementData displacementData =
                mapping::ElementDisplacementDataFromElementVDofs(
                    displacementElement, elementDisplacement
                );

            const mapping::ElementCompactificationData compactificationData(
                compactificationElement, elementCompactification
            );

            const mapping::ElementMappingData mappingData{
                .displacement     = displacementData,
                .compactification = compactificationData
            };

            const mfem::IntegrationRule &integrationRule = get_eos_rule(
                m_fem, m_barotrope, densityElement, enthalpyElement,
                *transformation
            );

            const int quadraturePointCount = integrationRule.GetNPoints();

            const int densityDofCount      = densityElement.GetDof();

            const int enthalpyDofCount     = enthalpyElement.GetDof();

            data.densityBasis.SetSize(quadraturePointCount, densityDofCount);

            data.enthalpyBasis.SetSize(quadraturePointCount, enthalpyDofCount);

            data.weightedResidual.SetSize(quadraturePointCount);

            data.quadratureWeights.SetSize(quadraturePointCount);

            data.weightedEnthalpyDerivative.SetSize(quadraturePointCount);

            densityShape.SetSize(densityDofCount);

            enthalpyShape.SetSize(enthalpyDofCount);

            for (int quadraturePoint = 0;
                 quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint =
                    integrationRule.IntPoint(quadraturePoint);

                transformation->SetIntPoint(&integrationPoint);

                mapping::VolumeMappingContext mappingContext;

                const mapping::MappingStatus mappingStatus =
                    m_domainMapper.EvaluateVolume(
                        mappingData, *transformation, integrationPoint,
                        workspace, mappingContext
                    );

                MFEM_VERIFY(
                    mappingStatus == mapping::MappingStatus::valid,
                    "Stateless mapping failed while preparing "
                    "the barotropic closure operator. Element: "
                        << elementId
                        << ", attribute: " << transformation->Attribute
                        << ", quadrature point: " << quadraturePoint
                        << ", status: " << static_cast<int>(mappingStatus)
                );

                densityElement.CalcShape(integrationPoint, densityShape);

                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                for (int densityDof = 0; densityDof < densityDofCount;
                     ++densityDof) {
                    data.densityBasis(quadraturePoint, densityDof) =
                        densityShape(densityDof);
                }

                for (int enthalpyDof = 0; enthalpyDof < enthalpyDofCount;
                     ++enthalpyDof) {
                    data.enthalpyBasis(quadraturePoint, enthalpyDof) =
                        enthalpyShape(enthalpyDof);
                }

                const double density  = elementBaseDensity * densityShape;

                const double enthalpy = elementBaseEnthalpy * enthalpyShape;

                const double quadratureWeight =
                    mappingContext.quadrature.weight;

                const double eosDensity =
                    m_barotrope.density_from_enthalpy(enthalpy);

                const double enthalpyDerivative =
                    m_barotrope.density_derivative_from_enthalpy(enthalpy);

                MFEM_VERIFY(
                    std::isfinite(quadratureWeight) && quadratureWeight > 0.0 &&
                        std::isfinite(eosDensity) &&
                        std::isfinite(enthalpyDerivative),
                    "PreparedBarotropicClosureOperator "
                    "encountered invalid quadrature data."
                );

                data.quadratureWeights(quadraturePoint) = quadratureWeight;

                data.weightedResidual(quadraturePoint) =
                    quadratureWeight * (density - eosDensity);

                data.weightedEnthalpyDerivative(quadraturePoint) =
                    quadratureWeight * enthalpyDerivative;
            }
        }

        MFEM_VERIFY(
            !m_elements.empty(), "PreparedBarotropicClosureOperator found no "
                                 "stellar elements."
        );

        m_baseDensityTrue      = baseDensityTrue;
        m_baseEnthalpyTrue     = baseEnthalpyTrue;
        m_baseDisplacementTrue = displacementTrue;

        m_isPrepared           = true;
        ++m_preparationCount;
    }

    void PreparedBarotropicClosureOperator::BuildResidual(
        mfem::Vector &residual
    ) const {
        MFEM_VERIFY(
            m_isPrepared, "PreparedBarotropicClosureOperator must be "
                          "prepared before BuildResidual is called."
        );

        mfem::Vector localResidual(m_fem.densityFes->GetVSize());
        localResidual = 0.0;

        mfem::Vector elementResidual;

        for (const ElementPAData &data : m_elements) {
            elementResidual.SetSize(data.densityDofs.Size());

            data.densityBasis.MultTranspose(
                data.weightedResidual, elementResidual
            );

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->TransformDual(elementResidual);
            }

            localResidual.AddElementVector(data.densityDofs, elementResidual);
        }

        local_to_true(*m_fem.densityFes, localResidual, residual);
    }

    void PreparedBarotropicClosureOperator::Mult(
        const mfem::Vector &densityVariationTrue,
        const mfem::Vector &enthalpyVariationTrue,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        MFEM_VERIFY(
            densityVariationTrue.Size() == m_densitySize,
            "The density-variation true vector has "
            "the wrong size."
        );

        MFEM_VERIFY(
            enthalpyVariationTrue.Size() == m_enthalpySize,
            "The enthalpy-variation true vector has "
            "the wrong size."
        );

        MFEM_VERIFY(
            displacementVariationTrue.Size() ==
                m_fem.displacementFes->GetTrueVSize(),
            "The displacement-variation true vector has "
            "the wrong size."
        );

        Mult(densityVariationTrue, enthalpyVariationTrue, action);

        mfem::Vector displacementAction;

        kernels::apply_barotropic_closure_displacement_action(
            m_fem, m_domainMapper, m_barotrope, m_baseDensityTrue,
            m_baseEnthalpyTrue, m_baseDisplacementTrue,
            displacementVariationTrue, displacementAction
        );

        MFEM_VERIFY(
            displacementAction.Size() == m_densitySize,
            "The barotropic-closure displacement action "
            "returned a vector with the wrong size."
        );

        action += displacementAction;
    }

    void PreparedBarotropicClosureOperator::Mult(
        const mfem::Vector &combinedVariation,
        mfem::Vector &action
    ) const {
        VerifyPrepared();

        const int displacementSize = m_fem.displacementFes->GetTrueVSize();

        const int combinedSize =
            m_densitySize + m_enthalpySize + displacementSize;

        MFEM_VERIFY(
            combinedVariation.Size() == combinedSize,
            "The combined barotropic-closure variation "
            "vector has the wrong size. Expected "
                << combinedSize << " entries but received "
                << combinedVariation.Size() << "."
        );

        mfem::real_t *combinedData =
            const_cast<mfem::real_t *>(combinedVariation.HostRead());

        const mfem::Vector densityVariationTrue(combinedData, m_densitySize);

        const mfem::Vector enthalpyVariationTrue(
            combinedData + m_densitySize, m_enthalpySize
        );

        const mfem::Vector displacementVariationTrue(
            combinedData + m_densitySize + m_enthalpySize, displacementSize
        );

        Mult(
            densityVariationTrue, enthalpyVariationTrue,
            displacementVariationTrue, action
        );
    }

    void PreparedBarotropicClosureOperator::Mult(
        const mfem::Vector &densityVariationTrue,
        const mfem::Vector &enthalpyVariationTrue,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_isPrepared, "PreparedBarotropicClosureOperator must be "
                          "prepared before Mult is called."
        );

        MFEM_VERIFY(
            densityVariationTrue.Size() == m_densitySize,
            "PreparedBarotropicClosureOperator received a "
            "density variation with the wrong size."
        );

        MFEM_VERIFY(
            enthalpyVariationTrue.Size() == m_enthalpySize,
            "PreparedBarotropicClosureOperator received an "
            "enthalpy variation with the wrong size."
        );

        mfem::Vector densityVariationLocal;
        mfem::Vector enthalpyVariationLocal;

        true_to_local(
            *m_fem.densityFes, densityVariationTrue, densityVariationLocal
        );

        true_to_local(
            *m_fem.enthalpyFes, enthalpyVariationTrue, enthalpyVariationLocal
        );

        mfem::Vector localAction(m_fem.densityFes->GetVSize());
        localAction = 0.0;

        mfem::Vector elementDensityVariation;
        mfem::Vector elementEnthalpyVariation;
        mfem::Vector quadratureDensityVariation;
        mfem::Vector quadratureEnthalpyVariation;
        mfem::Vector quadratureAction;
        mfem::Vector elementAction;

        for (const ElementPAData &data : m_elements) {
            densityVariationLocal.GetSubVector(
                data.densityDofs, elementDensityVariation
            );

            enthalpyVariationLocal.GetSubVector(
                data.enthalpyDofs, elementEnthalpyVariation
            );

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->InvTransformPrimal(
                    elementDensityVariation
                );
            }

            if (data.enthalpyDofTransformation != nullptr) {
                data.enthalpyDofTransformation->InvTransformPrimal(
                    elementEnthalpyVariation
                );
            }

            quadratureDensityVariation.SetSize(data.quadratureWeights.Size());

            quadratureEnthalpyVariation.SetSize(data.quadratureWeights.Size());

            quadratureAction.SetSize(data.quadratureWeights.Size());

            data.densityBasis.Mult(
                elementDensityVariation, quadratureDensityVariation
            );

            data.enthalpyBasis.Mult(
                elementEnthalpyVariation, quadratureEnthalpyVariation
            );

            for (int quadraturePoint = 0;
                 quadraturePoint < quadratureAction.Size(); ++quadraturePoint) {
                quadratureAction(quadraturePoint) =
                    data.quadratureWeights(quadraturePoint) *
                        quadratureDensityVariation(quadraturePoint) -
                    data.weightedEnthalpyDerivative(quadraturePoint) *
                        quadratureEnthalpyVariation(quadraturePoint);
            }

            elementAction.SetSize(data.densityDofs.Size());

            data.densityBasis.MultTranspose(quadratureAction, elementAction);

            if (data.densityDofTransformation != nullptr) {
                data.densityDofTransformation->TransformDual(elementAction);
            }

            localAction.AddElementVector(data.densityDofs, elementAction);
        }

        local_to_true(*m_fem.densityFes, localAction, action);
    }

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

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

    int PreparedBarotropicClosureOperator::GetDensitySize() const noexcept {
        return m_densitySize;
    }

    int PreparedBarotropicClosureOperator::GetEnthalpySize() const noexcept {
        return m_enthalpySize;
    }

    void PreparedBarotropicClosureOperator::VerifyPrepared() const {
        MFEM_VERIFY(
            m_isPrepared, "PreparedBarotropicClosureOperator must be "
                          "prepared before this operation is called."
        );
    }
} // namespace mean_field::operators