module;

#include <array>
#include <cmath>
#include <limits>

#include <mfem.hpp>

module mean_field;

import :operators.kernels.pressure_force;

namespace {
    void true_to_local(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &trueVector,
        mfem::Vector &localVector
    ) {
        MFEM_VERIFY(
            trueVector.Size() == finiteElementSpace.GetTrueVSize(),
            "The pressure-force 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(),
            "The pressure-force 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;
        }
    }

    [[nodiscard]] int vector_dof_index(
        const mfem::Ordering::Type ordering,
        const int scalarDof,
        const int component,
        const int scalarDofCount,
        const int dimension
    ) {
        if (ordering == mfem::Ordering::byNODES) {
            return scalarDof + component * scalarDofCount;
        }

        if (ordering == mfem::Ordering::byVDIM) {
            return component + scalarDof * dimension;
        }
        MFEM_ABORT("The displacement space uses an unsupported ordering.");
        return -1;
    }

    [[nodiscard]] int get_pressure_extra_order(
        const mean_field::physics::PolytropicBarotrope &barotrope
    ) {
        /*
         * Pressure has the enthalpy dependence
         *
         *     P(h) proportional to h^(n + 1).
         *
         * The registered enthalpy operand already contributes one factor
         * of the enthalpy polynomial order. The remaining dynamic
         * contribution is therefore n times that order.
         */
        const double extraOrder =
            barotrope.polytropic_index() *
            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 pressure EOS effective polynomial order is invalid."
        );

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

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

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

        MFEM_VERIFY(
            displacementElement.GetOrder() ==
                mean_field::field::Displacement::Vector::familyOrder,
            "The pressure-force test element does not match the "
            "registered displacement field."
        );

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

        const mean_field::quadrature::MfemRule rule =
            f.quadratureFactory->get(query, transformation.GetGeometryType());

        MFEM_VERIFY(
            rule.integration_rule != nullptr,
            "The quadrature policy did not return a pressure-force "
            "integration rule."
        );

        return *rule.integration_rule;
    }

    void validate_inputs(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapperStateless &domainMapper,
        const mfem::Vector &enthalpyTrue,
        const mfem::Vector &displacementTrue
    ) {
        MFEM_VERIFY(
            f.mesh != nullptr, "The pressure-force kernel requires a mesh."
        );

        MFEM_VERIFY(
            f.enthalpyFes != nullptr,
            "The pressure-force kernel requires the enthalpy "
            "finite-element space."
        );

        MFEM_VERIFY(
            f.displacementFes != nullptr,
            "The pressure-force kernel requires the displacement "
            "finite-element space."
        );

        MFEM_VERIFY(
            f.compactificationFes != nullptr,
            "The pressure-force kernel requires the compactification "
            "finite-element space."
        );

        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr,
            "The pressure-force kernel requires the compactification "
            "coordinate."
        );

        MFEM_VERIFY(
            f.quadratureFactory != nullptr,
            "The pressure-force kernel requires the quadrature "
            "rule factory."
        );

        MFEM_VERIFY(
            enthalpyTrue.Size() == f.enthalpyFes->GetTrueVSize(),
            "The pressure-force enthalpy vector has the wrong size."
        );

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

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

        MFEM_VERIFY(
            f.displacementFes->GetVDim() == f.mesh->Dimension(),
            "The displacement vector dimension does not match the "
            "mesh dimension."
        );

        MFEM_VERIFY(
            f.displacementFes->GetOrdering() == mfem::Ordering::byNODES,
            "The pressure-force kernel requires the registered byNODES "
            "displacement ordering."
        );
    }
} // namespace

namespace mean_field::operators::kernels {
    void apply_pressure_force_residual(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const physics::PolytropicBarotrope &barotrope,
        const mfem::Vector &enthalpyTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &residualTrue
    ) {
        validate_inputs(f, domainMapper, enthalpyTrue, displacementTrue);

        mfem::Vector enthalpyLocal;
        mfem::Vector displacementLocal;

        true_to_local(*f.enthalpyFes, enthalpyTrue, enthalpyLocal);

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

        mfem::Vector localResidual(f.displacementFes->GetVSize());
        localResidual = 0.0;

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

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

        mfem::Vector elementEnthalpy;
        mfem::Vector elementDisplacement;
        mfem::Vector elementCompactification;
        mfem::Vector elementResidual;
        mfem::Vector enthalpyShape;

        mfem::DenseMatrix displacementDShapeReference;
        mfem::DenseMatrix displacementDShapePhysical;

        mapping::VolumeMappingContext mappingContext;

        const int dimension       = f.mesh->Dimension();
        const int vacuumAttribute = domainMapper.GetVacuumElementAttribute();

        const mfem::Ordering::Type displacementOrdering =
            f.displacementFes->GetOrdering();

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

            MFEM_VERIFY(
                transformation != nullptr,
                "The pressure-force kernel received a null element "
                "transformation."
            );

            /*
             * Skip vacuum before constructing or evaluating any mapping
             * data for the element.
             */
            if (transformation->Attribute == vacuumAttribute) {
                continue;
            }

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

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

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

            mfem::DofTransformation *enthalpyDofTransformation =
                f.enthalpyFes->GetElementDofs(elementId, enthalpyDofs);

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

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

            enthalpyLocal.GetSubVector(enthalpyDofs, elementEnthalpy);

            displacementLocal.GetSubVector(
                displacementDofs, elementDisplacement
            );

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

            if (enthalpyDofTransformation != nullptr) {
                enthalpyDofTransformation->InvTransformPrimal(elementEnthalpy);
            }

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

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

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

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

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

            const int scalarDisplacementDofCount = displacementElement.GetDof();

            MFEM_VERIFY(
                displacementDofs.Size() ==
                    scalarDisplacementDofCount * dimension,
                "The pressure-force element displacement vector has "
                "the wrong size."
            );

            enthalpyShape.SetSize(enthalpyElement.GetDof());

            displacementDShapeReference.SetSize(
                scalarDisplacementDofCount, dimension
            );

            displacementDShapePhysical.SetSize(
                scalarDisplacementDofCount, dimension
            );

            elementResidual.SetSize(displacementDofs.Size());
            elementResidual = 0.0;

            const mfem::IntegrationRule &integrationRule =
                get_pressure_force_rule(
                    f, barotrope, enthalpyElement, displacementElement,
                    *transformation
                );

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

                transformation->SetIntPoint(&integrationPoint);

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

                MFEM_VERIFY(
                    mappingStatus == mapping::MappingStatus::valid,
                    "Stateless mapping failed in the pressure-force "
                    "kernel. Element: "
                        << elementId
                        << ", attribute: " << transformation->Attribute
                        << ", quadrature point: " << quadratureIndex
                        << ", status: " << static_cast<int>(mappingStatus)
                );

                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                const double enthalpyValue = elementEnthalpy * enthalpyShape;

                const double pressureValue =
                    barotrope.pressure_from_enthalpy(enthalpyValue);

                displacementElement.CalcDShape(
                    integrationPoint, displacementDShapeReference
                );

                /*
                 * Row i of DShape is grad_reference(N_i). Multiplication
                 * by the complete inverse element Jacobian gives
                 *
                 *     grad_physical(N_i)
                 *       = grad_reference(N_i) J^{-1}.
                 */
                mfem::Mult(
                    displacementDShapeReference,
                    mappingContext.quadrature.J_inv, displacementDShapePhysical
                );

                const double weightedPressure =
                    pressureValue * mappingContext.quadrature.weight;

                MFEM_VERIFY(
                    std::isfinite(pressureValue) &&
                        std::isfinite(weightedPressure),
                    "The pressure-force kernel encountered a non-finite "
                    "quadrature value."
                );

                /*
                 * For the vector basis N_i e_c,
                 *
                 *     div(N_i e_c) = partial_c N_i.
                 *
                 * Therefore
                 *
                 *     R_(i,c)
                 *       = -integral P partial_c N_i dV.
                 */
                for (int scalarDof = 0; scalarDof < scalarDisplacementDofCount;
                     ++scalarDof) {
                    for (int component = 0; component < dimension;
                         ++component) {
                        const int vectorDof = vector_dof_index(
                            displacementOrdering, scalarDof, component,
                            scalarDisplacementDofCount, dimension
                        );

                        elementResidual(vectorDof) -=
                            weightedPressure *
                            displacementDShapePhysical(scalarDof, component);
                    }
                }
            }

            if (displacementDofTransformation != nullptr) {
                displacementDofTransformation->TransformDual(elementResidual);
            }

            localResidual.AddElementVector(displacementDofs, elementResidual);
        }

        local_to_true(*f.displacementFes, localResidual, residualTrue);
    }
} // namespace mean_field::operators::kernels