module;

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

module mean_field;

import :operators.kernels.barotropic_closure;
import :field.registry;
import :utils.domain;

namespace {
    namespace eos       = mean_field::eos;

    using DomainSchema  = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using ClosureDomain = mean_field::field::FieldDomainT<mean_field::field::Density>;

    enum class ClosureAction { residual, density, enthalpy };

    [[nodiscard]] bool element_is_in_closure_support(const int attribute) {
        return DomainSchema::template attribute_belongs_to<ClosureDomain>(attribute);
    }

    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::eos::Polytrope &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::eos::Polytrope &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 EOS test element does not match the "
            "registered density field."
        );

        MFEM_VERIFY(
            enthalpyElement.GetOrder() == mean_field::field::Enthalpy::Scalar::familyOrder,
            "The 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 an "
                                                    "EOS-closure integration rule."
        );

        return *resolution.integration_rule;
    }

    void validate_common_inputs(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapper &domainMapper,
        const mfem::Vector &displacementTrue
    ) {
        MFEM_VERIFY(f.mesh != nullptr, "The EOS closure kernel requires a mesh.");
        MFEM_VERIFY(
            f.densityFes != nullptr, "The EOS closure kernel requires the density "
                                     "finite-element space."
        );
        MFEM_VERIFY(
            f.enthalpyFes != nullptr, "The EOS closure kernel requires the enthalpy "
                                      "finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "The EOS closure kernel requires the displacement "
                                          "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "The EOS closure kernel requires the "
                                              "compactification finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "The EOS closure kernel requires the "
                                                     "compactification coordinate."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "The EOS closure kernel requires the quadrature "
                                            "rule factory."
        );
        MFEM_VERIFY(
            displacementTrue.Size() == f.displacementFes->GetTrueVSize(), "The displacement vector has the wrong size."
        );
        MFEM_VERIFY(
            domainMapper.GetDimension() == f.mesh->Dimension(), "The domain-mapper dimension does not match "
                                                                "the mesh dimension."
        );
    }

    void apply_closure_action(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapper &domainMapper,
        const mean_field::eos::Polytrope &barotrope,
        const ClosureAction closureAction,
        const mfem::Vector *densityInputTrue,
        const mfem::Vector *baseEnthalpyTrue,
        const mfem::Vector *enthalpyVariationTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        validate_common_inputs(f, domainMapper, displacementTrue);

        if (closureAction == ClosureAction::residual || closureAction == ClosureAction::density) {
            MFEM_VERIFY(
                densityInputTrue != nullptr && densityInputTrue->Size() == f.densityFes->GetTrueVSize(),
                "The density input has the wrong size."
            );
        }

        if (closureAction == ClosureAction::residual || closureAction == ClosureAction::enthalpy) {
            MFEM_VERIFY(
                baseEnthalpyTrue != nullptr && baseEnthalpyTrue->Size() == f.enthalpyFes->GetTrueVSize(),
                "The base enthalpy has the wrong size."
            );
        }

        if (closureAction == ClosureAction::enthalpy) {
            MFEM_VERIFY(
                enthalpyVariationTrue != nullptr && enthalpyVariationTrue->Size() == f.enthalpyFes->GetTrueVSize(),
                "The enthalpy variation has the wrong size."
            );
        }

        mfem::Vector densityInputLocal;
        mfem::Vector baseEnthalpyLocal;
        mfem::Vector enthalpyVariationLocal;
        mfem::Vector displacementLocal;

        if (densityInputTrue != nullptr) {
            true_to_local(*f.densityFes, *densityInputTrue, densityInputLocal);
        }

        if (baseEnthalpyTrue != nullptr) {
            true_to_local(*f.enthalpyFes, *baseEnthalpyTrue, baseEnthalpyLocal);
        }

        if (enthalpyVariationTrue != nullptr) {
            true_to_local(*f.enthalpyFes, *enthalpyVariationTrue, enthalpyVariationLocal);
        }

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

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

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

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

        mfem::Vector elementDensityInput;
        mfem::Vector elementBaseEnthalpy;
        mfem::Vector elementEnthalpyVariation;
        mfem::Vector elementDisplacement;
        mfem::Vector elementCompactification;
        mfem::Vector elementAction;

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

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

            MFEM_VERIFY(
                transformation != nullptr, "The EOS closure kernel received a null "
                                           "element transformation."
            );

            if (!element_is_in_closure_support(transformation->Attribute)) {
                continue;
            }

            const mfem::FiniteElement &densityElement          = *f.densityFes->GetFE(elementId);
            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 *densityDofTransformation  = f.densityFes->GetElementDofs(elementId, densityDofs);

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

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

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

            if (densityInputTrue != nullptr) {
                densityInputLocal.GetSubVector(densityDofs, elementDensityInput);

                if (densityDofTransformation != nullptr) {
                    densityDofTransformation->InvTransformPrimal(elementDensityInput);
                }
            }

            if (baseEnthalpyTrue != nullptr) {
                baseEnthalpyLocal.GetSubVector(enthalpyDofs, elementBaseEnthalpy);

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

            if (enthalpyVariationTrue != nullptr) {
                enthalpyVariationLocal.GetSubVector(enthalpyDofs, elementEnthalpyVariation);

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

            displacementLocal.GetSubVector(displacementDofs, elementDisplacement);

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

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

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

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

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

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

            densityShape.SetSize(densityElement.GetDof());
            enthalpyShape.SetSize(enthalpyElement.GetDof());
            elementAction.SetSize(densityElement.GetDof());
            elementAction = 0.0;

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

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

                transformation->SetIntPoint(&integrationPoint);

                mean_field::mapping::VolumeMappingContext mappingContext;

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

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

                densityElement.CalcShape(integrationPoint, densityShape);

                double integrand = 0.0;

                if (closureAction == ClosureAction::density) {
                    integrand = elementDensityInput * densityShape;
                } else {
                    enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                    const double baseEnthalpy = elementBaseEnthalpy * enthalpyShape;

                    if (closureAction == ClosureAction::residual) {
                        const double density = elementDensityInput * densityShape;

                        const double equationOfStateDensity =
                            eos::evaluate<eos::quantity::Density>(barotrope, eos::SpecificEnthalpyValue{baseEnthalpy})
                                .value();

                        integrand = density - equationOfStateDensity;
                    } else {
                        const double enthalpyVariation = elementEnthalpyVariation * enthalpyShape;

                        const double densityDerivative =
                            eos::partialDerivative<eos::quantity::Density, eos::quantity::SpecificEnthalpy>(
                                barotrope, eos::SpecificEnthalpyValue{baseEnthalpy}
                            )
                                .value();

                        integrand = -densityDerivative * enthalpyVariation;
                    }
                }

                const double weightedIntegrand = mappingContext.quadrature.weight * integrand;

                for (int densityDof = 0; densityDof < densityElement.GetDof(); ++densityDof) {
                    elementAction(densityDof) += weightedIntegrand * densityShape(densityDof);
                }
            }

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

            localAction.AddElementVector(densityDofs, elementAction);
        }

        local_to_true(*f.densityFes, localAction, action);
    }
} // namespace

namespace mean_field::operators::kernels {
    void apply_barotropic_closure(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &barotrope,
        const mfem::Vector &densityTrue,
        const mfem::Vector &enthalpyTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &residual
    ) {
        apply_closure_action(
            f, domainMapper, barotrope, ClosureAction::residual, &densityTrue, &enthalpyTrue, nullptr, displacementTrue,
            residual
        );
    }

    void apply_barotropic_closure_density_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &barotrope,
        const mfem::Vector &densityVariationTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        apply_closure_action(
            f, domainMapper, barotrope, ClosureAction::density, &densityVariationTrue, nullptr, nullptr,
            displacementTrue, action
        );
    }

    void apply_barotropic_closure_enthalpy_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &barotrope,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &enthalpyVariationTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        apply_closure_action(
            f, domainMapper, barotrope, ClosureAction::enthalpy, nullptr, &baseEnthalpyTrue, &enthalpyVariationTrue,
            displacementTrue, action
        );
    }

    void apply_barotropic_closure_displacement_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const eos::Polytrope &barotrope,
        const mfem::Vector &baseDensityTrue,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &displacementTrue,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &action
    ) {
        MFEM_VERIFY(
            f.mesh != nullptr, "The barotropic-closure displacement action "
                               "requires a mesh."
        );

        MFEM_VERIFY(
            f.densityFes != nullptr, "The barotropic-closure displacement action "
                                     "requires the density finite-element space."
        );

        MFEM_VERIFY(
            f.enthalpyFes != nullptr, "The barotropic-closure displacement action "
                                      "requires the enthalpy finite-element space."
        );

        MFEM_VERIFY(
            f.displacementFes != nullptr, "The barotropic-closure displacement action "
                                          "requires the displacement finite-element space."
        );

        MFEM_VERIFY(
            f.compactificationFes != nullptr, "The barotropic-closure displacement action "
                                              "requires the compactification finite-element space."
        );

        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "The barotropic-closure displacement action "
                                                     "requires the compactification coordinate."
        );

        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "The barotropic-closure displacement action "
                                            "requires the quadrature-rule factory."
        );

        MFEM_VERIFY(
            baseDensityTrue.Size() == f.densityFes->GetTrueVSize(), "The base-density vector has the wrong size."
        );

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

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

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

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

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

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

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

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

        true_to_local(*f.displacementFes, displacementVariationTrue, displacementVariationLocal);

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

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

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

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

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

        mapping::VolumeMappingContext mappingContext;
        mapping::VolumeMappingVariation mappingVariation;

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

            MFEM_VERIFY(
                transformation != nullptr, "The barotropic-closure displacement action "
                                           "received a null element transformation."
            );

            if (!element_is_in_closure_support(transformation->Attribute)) {
                continue;
            }

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

            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 *densityDofTransformation  = f.densityFes->GetElementDofs(elementId, densityDofs);

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

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

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

            baseDensityLocal.GetSubVector(densityDofs, elementBaseDensity);

            baseEnthalpyLocal.GetSubVector(enthalpyDofs, elementBaseEnthalpy);

            displacementLocal.GetSubVector(displacementDofs, elementDisplacement);

            displacementVariationLocal.GetSubVector(displacementDofs, elementDisplacementVariation);

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

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

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

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

                displacementDofTransformation->InvTransformPrimal(elementDisplacementVariation);
            }

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

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

            const mapping::ElementDisplacementData displacementVariationData =
                mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementDisplacementVariation);

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

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

            densityShape.SetSize(densityElement.GetDof());

            enthalpyShape.SetSize(enthalpyElement.GetDof());

            elementAction.SetSize(densityElement.GetDof());
            elementAction = 0.0;

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

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

                transformation->SetIntPoint(&integrationPoint);

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

                MFEM_VERIFY(
                    mappingStatus == mapping::MappingStatus::valid,
                    "The base mapping is invalid while applying "
                    "the barotropic-closure displacement action. "
                    "Element: "
                        << elementId << ", attribute: " << transformation->Attribute
                        << ", quadrature point: " << quadraturePoint << ", status: " << static_cast<int>(mappingStatus)
                );

                const mapping::MappingStatus variationStatus = domainMapper.EvaluateVolumeVariation(
                    mappingData, displacementVariationData, *transformation, integrationPoint, mappingContext,
                    workspace, mappingVariation
                );

                MFEM_VERIFY(
                    variationStatus == mapping::MappingStatus::valid,
                    "The mapping variation is invalid while "
                    "applying the barotropic-closure "
                    "displacement action. Element: "
                        << elementId << ", attribute: " << transformation->Attribute << ", quadrature point: "
                        << quadraturePoint << ", status: " << static_cast<int>(variationStatus)
                );

                densityElement.CalcShape(integrationPoint, densityShape);

                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                const double densityValue  = elementBaseDensity * densityShape;

                const double enthalpyValue = elementBaseEnthalpy * enthalpyShape;

                const double equationOfStateDensity =
                    eos::evaluate<eos::quantity::Density>(barotrope, eos::SpecificEnthalpyValue{enthalpyValue}).value();

                const double closureValue        = densityValue - equationOfStateDensity;

                const double geometryActionValue = closureValue * mappingVariation.weight_variation;

                MFEM_VERIFY(
                    std::isfinite(closureValue) && std::isfinite(geometryActionValue),
                    "The barotropic-closure displacement action "
                    "encountered a non-finite quadrature value."
                );

                elementAction.Add(geometryActionValue, densityShape);
            }

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

            localAction.AddElementVector(densityDofs, elementAction);
        }

        local_to_true(*f.densityFes, localAction, action);
    }
} // namespace mean_field::operators::kernels