module;

#include <algorithm>
#include <array>
#include <cmath>
#include <mfem.hpp>
#include <optional>

module mean_field;

import :operators.kernels.hydrostatic_equilibrium;

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

    [[nodiscard]] bool is_vacuum_attribute(const int attribute) {
        return DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Vacuum>(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;
        }
    }

    void validate_fem(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapper &domainMapper
    ) {
        MFEM_VERIFY(f.mesh != nullptr, "The hydrostatic kernel requires a mesh.");

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

        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "The hydrostatic kernel requires the "
                                              "gravity-potential finite-element space."
        );

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

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

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

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

        MFEM_VERIFY(
            f.mesh->Dimension() == 3, "The rigid-rotation hydrostatic kernel "
                                      "currently requires a three-dimensional mesh."
        );

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

    const mfem::IntegrationRule &get_hydrostatic_rule(
        const mean_field::fem::FEM &f,
        const mfem::FiniteElement &enthalpyElement,
        const mfem::FiniteElement &potentialElement,
        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 hydrostatic test element does not match "
            "the registered enthalpy field."
        );

        MFEM_VERIFY(
            potentialElement.GetOrder() == mean_field::field::Gravity::Potential::familyOrder,
            "The hydrostatic potential element does not "
            "match the registered gravity-potential field."
        );

        const auto enthalpyQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumEnthalpy>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        const auto gravityQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumGravity>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        const auto rotationQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumRotation>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), std::array<int, 1>{2},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        const auto constantQuery = EnthalpyField::make_query<mean_field::field::Enthalpy::Form::EquilibriumConstant>(
            mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
            mean_field::utils::DOMAINS::STELLAR, mean_field::quadrature::MappingKind::general
        );

        int integrationOrder    = 0;

        const auto update_order = [&f, &transformation, &integrationOrder](const mean_field::quadrature::Query &query) {
            const auto rule = f.quadratureFactory->get(query, transformation.GetGeometryType());

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

            integrationOrder = std::max(integrationOrder, rule.resolution.order);
        };

        update_order(enthalpyQuery);
        update_order(gravityQuery);
        update_order(rotationQuery);
        update_order(constantQuery);

        return mfem::IntRules.Get(transformation.GetGeometryType(), integrationOrder);
    }

    struct HydrostaticAssemblyRequest {
        const mean_field::physics::RigidRotation *rotation{nullptr};

        const mfem::Vector *baseEnthalpyTrue{nullptr};
        const mfem::Vector *basePotentialTrue{nullptr};

        const mfem::Vector *enthalpyVariationTrue{nullptr};
        const mfem::Vector *potentialVariationTrue{nullptr};
        const mfem::Vector *displacementVariationTrue{nullptr};

        double bernoulliConstant{0.0};
        double constantVariation{0.0};

        bool buildResidual{false};
    };

    void assemble_hydrostatic_form(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapper &domainMapper,
        const mfem::Vector &displacementTrue,
        const HydrostaticAssemblyRequest &request,
        mfem::Vector &result
    ) {
        validate_fem(f, domainMapper);

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

        MFEM_VERIFY(std::isfinite(request.bernoulliConstant), "The Bernoulli constant is non-finite.");

        MFEM_VERIFY(std::isfinite(request.constantVariation), "The Bernoulli-constant variation is non-finite.");

        const bool requiresBaseState = request.buildResidual || request.displacementVariationTrue != nullptr;

        if (requiresBaseState) {
            MFEM_VERIFY(
                request.rotation != nullptr, "The hydrostatic residual or geometry "
                                             "action requires the rotation model."
            );

            MFEM_VERIFY(
                request.baseEnthalpyTrue != nullptr, "The hydrostatic residual or geometry "
                                                     "action requires the base enthalpy."
            );

            MFEM_VERIFY(
                request.basePotentialTrue != nullptr, "The hydrostatic residual or geometry "
                                                      "action requires the base potential."
            );
        }

        if (request.baseEnthalpyTrue != nullptr) {
            MFEM_VERIFY(
                request.baseEnthalpyTrue->Size() == f.enthalpyFes->GetTrueVSize(),
                "The base enthalpy vector has the wrong size."
            );
        }

        if (request.basePotentialTrue != nullptr) {
            MFEM_VERIFY(
                request.basePotentialTrue->Size() == f.gravityPotentialFes->GetTrueVSize(),
                "The base potential vector has the wrong size."
            );
        }

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

        if (request.potentialVariationTrue != nullptr) {
            MFEM_VERIFY(
                request.potentialVariationTrue->Size() == f.gravityPotentialFes->GetTrueVSize(),
                "The potential variation has the wrong size."
            );
        }

        if (request.displacementVariationTrue != nullptr) {
            MFEM_VERIFY(
                request.displacementVariationTrue->Size() == f.displacementFes->GetTrueVSize(),
                "The displacement variation has the wrong size."
            );
        }

        mfem::Vector displacementLocal;
        true_to_local(*f.displacementFes, displacementTrue, displacementLocal);

        mfem::Vector baseEnthalpyLocal;
        mfem::Vector basePotentialLocal;
        mfem::Vector enthalpyVariationLocal;
        mfem::Vector potentialVariationLocal;
        mfem::Vector displacementVariationLocal;

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

        if (request.basePotentialTrue != nullptr) {
            true_to_local(*f.gravityPotentialFes, *request.basePotentialTrue, basePotentialLocal);
        }

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

        if (request.potentialVariationTrue != nullptr) {
            true_to_local(*f.gravityPotentialFes, *request.potentialVariationTrue, potentialVariationLocal);
        }

        if (request.displacementVariationTrue != nullptr) {
            true_to_local(*f.displacementFes, *request.displacementVariationTrue, displacementVariationLocal);
        }

        mfem::Vector localResult(f.enthalpyFes->GetVSize());

        localResult = 0.0;

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

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

        mfem::Vector elementBaseEnthalpy;
        mfem::Vector elementBasePotential;
        mfem::Vector elementEnthalpyVariation;
        mfem::Vector elementPotentialVariation;
        mfem::Vector elementDisplacement;
        mfem::Vector elementDisplacementVariation;
        mfem::Vector elementCompactification;
        mfem::Vector elementResult;

        mfem::Vector enthalpyShape;
        mfem::Vector potentialShape;

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

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

            if (is_vacuum_attribute(transformation->Attribute)) {
                continue;
            }

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

            const mfem::FiniteElement &potentialElement        = *f.gravityPotentialFes->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 *potentialDofTransformation =
                f.gravityPotentialFes->GetElementDofs(elementId, potentialDofs);

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

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

            displacementLocal.GetSubVector(displacementDofs, elementDisplacement);

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

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

            if (request.basePotentialTrue != nullptr) {
                basePotentialLocal.GetSubVector(potentialDofs, elementBasePotential);
            }

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

            if (request.potentialVariationTrue != nullptr) {
                potentialVariationLocal.GetSubVector(potentialDofs, elementPotentialVariation);
            }

            if (request.displacementVariationTrue != nullptr) {
                displacementVariationLocal.GetSubVector(displacementDofs, elementDisplacementVariation);
            }

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

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

            if (potentialDofTransformation != nullptr) {
                if (request.basePotentialTrue != nullptr) {
                    potentialDofTransformation->InvTransformPrimal(elementBasePotential);
                }

                if (request.potentialVariationTrue != nullptr) {
                    potentialDofTransformation->InvTransformPrimal(elementPotentialVariation);
                }
            }

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

                if (request.displacementVariationTrue != nullptr) {
                    displacementDofTransformation->InvTransformPrimal(elementDisplacementVariation);
                }
            }

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

            std::optional<mean_field::mapping::ElementDisplacementData> displacementVariationData;

            if (request.displacementVariationTrue != nullptr) {
                displacementVariationData.emplace(
                    mean_field::mapping::ElementDisplacementDataFromElementVDofs(
                        displacementElement, elementDisplacementVariation
                    )
                );
            }

            elementResult.SetSize(enthalpyElement.GetDof());

            elementResult = 0.0;

            enthalpyShape.SetSize(enthalpyElement.GetDof());

            potentialShape.SetSize(potentialElement.GetDof());

            const mfem::IntegrationRule &integrationRule =
                get_hydrostatic_rule(f, enthalpyElement, potentialElement, *transformation);

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

                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,
                    "The base mapping is invalid in the "
                    "hydrostatic kernel. Element: "
                        << elementId << ", quadrature point: " << quadraturePoint
                        << ", status: " << static_cast<int>(mappingStatus)
                );

                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                potentialElement.CalcShape(integrationPoint, potentialShape);

                double baseIntegrand = 0.0;

                if (requiresBaseState) {
                    const double enthalpyValue  = elementBaseEnthalpy * enthalpyShape;

                    const double potentialValue = elementBasePotential * potentialShape;

                    const double rotationPotential =
                        request.rotation->potential(mappingContext.mapping.physical_position);

                    baseIntegrand = enthalpyValue + potentialValue - rotationPotential - request.bernoulliConstant;
                }

                if (request.buildResidual) {
                    elementResult.Add(mappingContext.quadrature.weight * baseIntegrand, enthalpyShape);

                    continue;
                }

                double materialVariation = -request.constantVariation;

                if (request.enthalpyVariationTrue != nullptr) {
                    materialVariation += elementEnthalpyVariation * enthalpyShape;
                }

                if (request.potentialVariationTrue != nullptr) {
                    materialVariation += elementPotentialVariation * potentialShape;
                }

                double weightedVariation = mappingContext.quadrature.weight * materialVariation;

                if (request.displacementVariationTrue != nullptr) {
                    mean_field::mapping::VolumeMappingVariation mappingVariation;

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

                    MFEM_VERIFY(
                        variationStatus == mean_field::mapping::MappingStatus::valid,
                        "The mapping variation is invalid "
                        "in the hydrostatic kernel."
                    );

                    const double rotationVariation = request.rotation->potential_directional_derivative(
                        mappingContext.mapping.physical_position, mappingVariation.mapping.physical_position_variation
                    );

                    weightedVariation += baseIntegrand * mappingVariation.weight_variation -
                                         rotationVariation * mappingContext.quadrature.weight;
                }

                elementResult.Add(weightedVariation, enthalpyShape);
            }

            if (enthalpyDofTransformation != nullptr) {
                enthalpyDofTransformation->TransformDual(elementResult);
            }

            localResult.AddElementVector(enthalpyDofs, elementResult);
        }

        local_to_true(*f.enthalpyFes, localResult, result);
    }
} // namespace

namespace mean_field::operators::kernels {
    void apply_hydrostatic_equilibrium(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const physics::RigidRotation &rotation,
        const mfem::Vector &enthalpyTrue,
        const mfem::Vector &potentialTrue,
        const mfem::Vector &displacementTrue,
        const double bernoulliConstant,
        mfem::Vector &residual
    ) {
        HydrostaticAssemblyRequest request;

        request.rotation          = &rotation;
        request.baseEnthalpyTrue  = &enthalpyTrue;
        request.basePotentialTrue = &potentialTrue;
        request.bernoulliConstant = bernoulliConstant;
        request.buildResidual     = true;

        assemble_hydrostatic_form(f, domainMapper, displacementTrue, request, residual);
    }

    void apply_hydrostatic_equilibrium_enthalpy_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const mfem::Vector &enthalpyVariationTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.enthalpyVariationTrue = &enthalpyVariationTrue;

        assemble_hydrostatic_form(f, domainMapper, displacementTrue, request, action);
    }

    void apply_hydrostatic_equilibrium_potential_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const mfem::Vector &potentialVariationTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.potentialVariationTrue = &potentialVariationTrue;

        assemble_hydrostatic_form(f, domainMapper, displacementTrue, request, action);
    }

    void apply_hydrostatic_equilibrium_constant_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const double constantVariation,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.constantVariation = constantVariation;

        assemble_hydrostatic_form(f, domainMapper, displacementTrue, request, action);
    }

    void apply_hydrostatic_equilibrium_displacement_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const physics::RigidRotation &rotation,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &basePotentialTrue,
        const mfem::Vector &baseDisplacementTrue,
        const double baseBernoulliConstant,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.rotation                  = &rotation;
        request.baseEnthalpyTrue          = &baseEnthalpyTrue;
        request.basePotentialTrue         = &basePotentialTrue;
        request.displacementVariationTrue = &displacementVariationTrue;
        request.bernoulliConstant         = baseBernoulliConstant;

        assemble_hydrostatic_form(f, domainMapper, baseDisplacementTrue, request, action);
    }

    void apply_hydrostatic_equilibrium_action(
        const fem::FEM &f,
        const mapping::DomainMapper &domainMapper,
        const physics::RigidRotation &rotation,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &basePotentialTrue,
        const mfem::Vector &baseDisplacementTrue,
        const double baseBernoulliConstant,
        const mfem::Vector &enthalpyVariationTrue,
        const mfem::Vector &potentialVariationTrue,
        const double constantVariation,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.rotation                  = &rotation;
        request.baseEnthalpyTrue          = &baseEnthalpyTrue;
        request.basePotentialTrue         = &basePotentialTrue;
        request.enthalpyVariationTrue     = &enthalpyVariationTrue;
        request.potentialVariationTrue    = &potentialVariationTrue;
        request.displacementVariationTrue = &displacementVariationTrue;
        request.bernoulliConstant         = baseBernoulliConstant;
        request.constantVariation         = constantVariation;

        assemble_hydrostatic_form(f, domainMapper, baseDisplacementTrue, request, action);
    }
} // namespace mean_field::operators::kernels