module;

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

#include <mfem.hpp>

module mean_field;

import :operators.kernels.pressure_force;

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

enum class PressureForceAction { residual, enthalpy, displacement };

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 scalarDof * dimension + component;
  }

  MFEM_ABORT("The displacement space uses an unsupported ordering.");
  return -1;
}

[[nodiscard]] int
get_pressure_extra_order(const mean_field::eos::Polytrope &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::eos::Polytrope &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::DomainMapper &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.");

  /*
   * ElementDisplacementDataFromElementVDofs currently consumes the
   * registered byNODES layout. Keep this explicit so a future
   * registry change fails immediately rather than silently
   * corrupting the geometry.
   */
  MFEM_VERIFY(f.displacementFes->GetOrdering() == mfem::Ordering::byNODES,
              "The pressure-force kernel requires the registered byNODES "
              "displacement ordering.");
}

void apply_pressure_force_action(
    const mean_field::fem::FEM &f,
    const mean_field::mapping::DomainMapper &domainMapper,
    const mean_field::eos::Polytrope &barotrope,
    const PressureForceAction pressureForceAction,
    const mfem::Vector &baseEnthalpyTrue,
    const mfem::Vector *enthalpyVariationTrue,
    const mfem::Vector *displacementVariationTrue,
    const mfem::Vector &displacementTrue, mfem::Vector &actionTrue) {
  validate_inputs(f, domainMapper, baseEnthalpyTrue, displacementTrue);

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

  if (pressureForceAction == PressureForceAction::displacement) {
    MFEM_VERIFY(displacementVariationTrue != nullptr &&
                    displacementVariationTrue->Size() ==
                        f.displacementFes->GetTrueVSize(),
                "The pressure-force displacement variation has the wrong "
                "size.");
  }

  mfem::Vector baseEnthalpyLocal;
  mfem::Vector enthalpyVariationLocal;
  mfem::Vector displacementLocal;
  mfem::Vector displacementVariationLocal;

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

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

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

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

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

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

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

  mfem::Vector elementBaseEnthalpy;
  mfem::Vector elementEnthalpyVariation;
  mfem::Vector elementDisplacement;
  mfem::Vector elementDisplacementVariation;
  mfem::Vector elementCompactification;
  mfem::Vector elementAction;
  mfem::Vector enthalpyShape;

  mfem::Array<int> enthalpyDofs;

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

  mean_field::mapping::VolumeMappingContext mappingContext;

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

  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 (is_vacuum_attribute(transformation->Attribute)) {
      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);

    baseEnthalpyLocal.GetSubVector(enthalpyDofs, elementBaseEnthalpy);

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

    displacementLocal.GetSubVector(displacementDofs, elementDisplacement);

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

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

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

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

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

      if (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 (displacementVariationTrue != nullptr) {
      displacementVariationData.emplace(
          mean_field::mapping::ElementDisplacementDataFromElementVDofs(
              displacementElement, elementDisplacementVariation));
    }

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

    displacementDShapePhysicalVariation.SetSize(scalarDisplacementDofCount,
                                                dimension);

    elementAction.SetSize(displacementDofs.Size());
    elementAction = 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 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 pressure-force "
                  "kernel. Element: "
                      << elementId
                      << ", attribute: " << transformation->Attribute
                      << ", quadrature point: " << quadratureIndex
                      << ", status: " << static_cast<int>(mappingStatus));

      enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

      const double enthalpyValue = elementBaseEnthalpy * enthalpyShape;

      double pressureFactor = 0.0;

      if (pressureForceAction == PressureForceAction::residual ||
          pressureForceAction == PressureForceAction::displacement) {
        pressureFactor = barotrope.pressure_from_enthalpy(enthalpyValue);
      } else {
        const double enthalpyVariationValue =
            elementEnthalpyVariation * enthalpyShape;

        pressureFactor =
            barotrope.pressure_derivative_from_enthalpy(enthalpyValue) *
            enthalpyVariationValue;
      }

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

      std::optional<mean_field::mapping::VolumeMappingVariation>
          mappingVariation;

      if (pressureForceAction == PressureForceAction::displacement) {
        mappingVariation.emplace();

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

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

        /*
         * Differentiating
         *
         *     grad_x(N_i) = grad_reference(N_i) J^{-1}
         *
         * at the frozen base geometry gives the physical
         * test-gradient variation used by the geometric
         * pressure block.
         */
        mfem::Mult(displacementDShapeReference,
                   mappingVariation->inverse_element_jacobian_variation,
                   displacementDShapePhysicalVariation);
      }

      const double weightedPressureFactor =
          pressureFactor * mappingContext.quadrature.weight;

      MFEM_VERIFY(std::isfinite(pressureFactor) &&
                      std::isfinite(weightedPressureFactor),
                  "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);

          if (pressureForceAction == PressureForceAction::displacement) {
            /*
             * Differentiate the complete discrete factor
             *
             *     grad_x(N_i) dV_x.
             *
             * The enthalpy DOFs, and therefore P(h), are
             * frozen in this Jacobian column.
             */
            const double gradientWeightVariation =
                mappingContext.quadrature.weight *
                    displacementDShapePhysicalVariation(scalarDof, component) +
                mappingVariation->weight_variation *
                    displacementDShapePhysical(scalarDof, component);

            const double contribution =
                pressureFactor * gradientWeightVariation;

            MFEM_VERIFY(std::isfinite(gradientWeightVariation) &&
                            std::isfinite(contribution),
                        "The pressure-force geometry action "
                        "encountered a non-finite contribution.");

            elementAction(vectorDof) -= contribution;
          } else {
            elementAction(vectorDof) -=
                weightedPressureFactor *
                displacementDShapePhysical(scalarDof, component);
          }
        }
      }
    }

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

    localAction.AddElementVector(displacementDofs, elementAction);
  }

  local_to_true(*f.displacementFes, localAction, actionTrue);
}
} // namespace

namespace mean_field::operators::kernels {
void apply_pressure_force_residual(
    const fem::FEM &f, const mapping::DomainMapper &domainMapper,
    const eos::Polytrope &barotrope, const mfem::Vector &enthalpyTrue,
    const mfem::Vector &displacementTrue, mfem::Vector &residualTrue) {
  apply_pressure_force_action(f, domainMapper, barotrope,
                              PressureForceAction::residual, enthalpyTrue,
                              nullptr, nullptr, displacementTrue, residualTrue);
}

void apply_pressure_force_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 &actionTrue) {
  apply_pressure_force_action(f, domainMapper, barotrope,
                              PressureForceAction::enthalpy, baseEnthalpyTrue,
                              &enthalpyVariationTrue, nullptr, displacementTrue,
                              actionTrue);
}

void apply_pressure_force_displacement_action(
    const fem::FEM &f, const mapping::DomainMapper &domainMapper,
    const eos::Polytrope &barotrope, const mfem::Vector &baseEnthalpyTrue,
    const mfem::Vector &displacementVariationTrue,
    const mfem::Vector &displacementTrue, mfem::Vector &actionTrue) {
  apply_pressure_force_action(
      f, domainMapper, barotrope, PressureForceAction::displacement,
      baseEnthalpyTrue, nullptr, &displacementVariationTrue, displacementTrue,
      actionTrue);
}
} // namespace mean_field::operators::kernels