#include <algorithm>
#include <array>
#include <cmath>

#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace prepared_hydrostatic_analytic_solve_test_utils {
constexpr double bernoulliConstant = 0.83;
constexpr double enthalpyAmplitude = 0.61;

struct AnalyticCase {
  const char *name;

  std::array<double, 3> deformationScale;
  std::array<double, 3> angularVelocity;
  std::array<double, 3> rotationCenter;
};

class EnthalpyJacobianOperator final : public mfem::Operator {
public:
  EnthalpyJacobianOperator(
      const int enthalpySize,
      const mean_field::operators::PreparedHydrostaticEquilibriumOperator
          &preparedOperator)
      : mfem::Operator(enthalpySize), m_preparedOperator(preparedOperator) {}

  void Mult(const mfem::Vector &direction,
            mfem::Vector &action) const override {
    m_preparedOperator.ApplyEnthalpyJacobianAction(direction, action);
  }

private:
  const mean_field::operators::PreparedHydrostaticEquilibriumOperator
      &m_preparedOperator;
};

mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies
make_dependencies() {
  return {.discretization = {.identity = 701, .revision = 2},
          .enthalpy = {.identity = 709, .revision = 3},
          .gravityPotential = {.identity = 719, .revision = 5},
          .displacement = {.identity = 727, .revision = 7},
          .rotation = {.identity = 733, .revision = 11},
          .bernoulliConstant = {.identity = 739, .revision = 13}};
}

mean_field::operators::context::hydrostatic::HydrostaticEquilibriumStateView
make_state(const mfem::Vector &enthalpy, const mfem::Vector &gravityPotential,
           const mfem::Vector &displacement) {
  return {.enthalpy = enthalpy,
          .gravityPotential = gravityPotential,
          .displacement = displacement,
          .bernoulliConstant = bernoulliConstant};
}

mfem::Vector make_vector(const std::array<double, 3> &values) {
  mfem::Vector vector(3);

  for (int component = 0; component < 3; ++component) {
    vector(component) = values[static_cast<std::size_t>(component)];
  }

  return vector;
}

mean_field::physics::RigidRotation
make_rotation(const AnalyticCase &analyticCase) {
  return mean_field::physics::RigidRotation(
      make_vector(analyticCase.angularVelocity),
      make_vector(analyticCase.rotationCenter));
}

void map_to_physical(const mfem::Vector &referencePosition,
                     const AnalyticCase &analyticCase,
                     mfem::Vector &physicalPosition) {
  physicalPosition.SetSize(3);

  for (int component = 0; component < 3; ++component) {
    physicalPosition(component) =
        analyticCase.deformationScale[static_cast<std::size_t>(component)] *
        referencePosition(component);
  }
}

double exact_enthalpy_value(const mfem::Vector &referencePosition) {
  double normalizedRadiusSquared = 0.0;

  for (int component = 0; component < 3; ++component) {
    const double normalizedCoordinate =
        referencePosition(component) / mean_field::utils::RADIUS;

    normalizedRadiusSquared += normalizedCoordinate * normalizedCoordinate;
  }

  return enthalpyAmplitude * std::max(0.0, 1.0 - normalizedRadiusSquared);
}

double
exact_potential_value(const mfem::Vector &referencePosition,
                      const AnalyticCase &analyticCase,
                      const mean_field::physics::RigidRotation &rotation) {
  mfem::Vector physicalPosition;

  map_to_physical(referencePosition, analyticCase, physicalPosition);

  /*
   * Construct Phi so that
   *
   *     h + Phi - Psi_rotation - C = 0
   *
   * analytically.
   */
  return bernoulliConstant + rotation.potential(physicalPosition) -
         exact_enthalpy_value(referencePosition);
}

mfem::Array<int> make_stellar_element_marker(const mean_field::fem::FEM &f) {
  mfem::Array<int> stellarElementMarker(f.mesh->GetNE());

  const int vacuumAttribute = field_dof_test_utils::vacuum_material_attribute;

  for (int elementId = 0; elementId < f.mesh->GetNE(); ++elementId) {
    stellarElementMarker[elementId] =
        f.mesh->GetAttribute(elementId) != vacuumAttribute;
  }

  return stellarElementMarker;
}
} // namespace prepared_hydrostatic_analytic_solve_test_utils

TEST_CASE("Prepared Hydrostatic Operator Solves Analytic Bernoulli Equilibria",
          tags::barotrope_hydrostatic_prepared_analytic &tags::convergence
              &tags::accuracy) {
  using prepared_hydrostatic_analytic_solve_test_utils::AnalyticCase;

  constexpr double deformationX = 1.08;
  constexpr double deformationY = 0.96;

  /*
   * The third scale makes the affine deformation
   * volume-preserving:
   *
   *     det(F) = sx * sy * sz = 1.
   */
  constexpr double deformationZ = 1.0 / (deformationX * deformationY);

  const std::array<AnalyticCase, 3> analyticCases{
      {{.name = "spherical nonrotating equilibrium",
        .deformationScale = {1.0, 1.0, 1.0},
        .angularVelocity = {0.0, 0.0, 0.0},
        .rotationCenter = {0.0, 0.0, 0.0}},
       {.name = "spherical rotating equilibrium",
        .deformationScale = {1.0, 1.0, 1.0},
        .angularVelocity = {0.13, -0.09, 0.31},
        .rotationCenter = {0.04, -0.03, 0.02}},
       {.name = "volume-preserving deformed rotating equilibrium",
        .deformationScale = {deformationX, deformationY, deformationZ},
        .angularVelocity = {0.17, -0.12, 0.43},
        .rotationCenter = {0.031, -0.024, 0.018}}}};

  auto args = test_utils::setup_args();

  mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

  const MPI_Comm communicator = f.mesh->GetComm();

  const mean_field::field::FieldDofMap enthalpyMap =
      field_dof_test_utils::make_map<mean_field::field::Enthalpy>(
          *f.enthalpyFes);

  const mean_field::field::FieldDofMap gravityPotentialMap =
      field_dof_test_utils::make_map<mean_field::field::Gravity>(
          *f.gravityPotentialFes);

  const mean_field::field::FieldDofMap displacementMap =
      field_dof_test_utils::make_map<mean_field::field::Displacement>(
          *f.displacementFes);

  const mfem::Array<int> stellarElementMarker =
      prepared_hydrostatic_analytic_solve_test_utils::
          make_stellar_element_marker(f);

  for (const AnalyticCase &analyticCase : analyticCases) {
    DYNAMIC_SECTION(analyticCase.name) {
      const double deformationDeterminant = analyticCase.deformationScale[0] *
                                            analyticCase.deformationScale[1] *
                                            analyticCase.deformationScale[2];

      REQUIRE(std::abs(deformationDeterminant - 1.0) < 2.0e-14);

      const mean_field::physics::RigidRotation rotation =
          prepared_hydrostatic_analytic_solve_test_utils::make_rotation(
              analyticCase);

      auto displacementFunction =
          [&analyticCase](const mfem::Vector &referencePosition,
                          mfem::Vector &displacementValue) {
            mfem::Vector physicalPosition;

            prepared_hydrostatic_analytic_solve_test_utils::map_to_physical(
                referencePosition, analyticCase, physicalPosition);

            displacementValue.SetSize(3);
            displacementValue = physicalPosition;
            displacementValue -= referencePosition;
          };

      auto potentialFunction = [&analyticCase, &rotation](
                                   const mfem::Vector &referencePosition) {
        return prepared_hydrostatic_analytic_solve_test_utils::
            exact_potential_value(referencePosition, analyticCase, rotation);
      };

      auto enthalpyFunction = [](const mfem::Vector &referencePosition) {
        return prepared_hydrostatic_analytic_solve_test_utils::
            exact_enthalpy_value(referencePosition);
      };

      mfem::VectorFunctionCoefficient displacementCoefficient(
          f.mesh->Dimension(), displacementFunction);

      mfem::FunctionCoefficient potentialCoefficient(potentialFunction);

      mfem::FunctionCoefficient exactEnthalpyCoefficient(enthalpyFunction);

      /*
       * Project the prescribed geometry and potential.
       */
      mfem::ParGridFunction displacementField(f.displacementFes.get());

      mfem::ParGridFunction potentialField(f.gravityPotentialFes.get());

      displacementField.ProjectCoefficient(displacementCoefficient);

      potentialField.ProjectCoefficient(potentialCoefficient);

      mfem::Vector displacementTrue;
      mfem::Vector gravityPotentialTrue;

      displacementField.GetTrueDofs(displacementTrue);
      potentialField.GetTrueDofs(gravityPotentialTrue);

      const mfem::Vector displacement =
          displacementMap.gather(displacementTrue);
      const mfem::Vector gravityPotential =
          gravityPotentialMap.gather(gravityPotentialTrue);

      /*
       * This projection is not used as the solution. It gives
       * the best directly available representation baseline
       * against which the solved field can be compared.
       */
      mfem::ParGridFunction projectedEnthalpyField(f.enthalpyFes.get());

      projectedEnthalpyField.ProjectCoefficient(exactEnthalpyCoefficient);

      mfem::ParGridFunction zeroEnthalpyField(f.enthalpyFes.get());

      zeroEnthalpyField = 0.0;

      const double exactEnthalpyNorm = zeroEnthalpyField.ComputeL2Error(
          exactEnthalpyCoefficient, nullptr, &stellarElementMarker);

      const double projectionError = projectedEnthalpyField.ComputeL2Error(
          exactEnthalpyCoefficient, nullptr, &stellarElementMarker);

      REQUIRE(exactEnthalpyNorm > 0.0);

      const double relativeProjectionError =
          projectionError / exactEnthalpyNorm;

      /*
       * Begin deliberately far from equilibrium.
       */
      mfem::Vector enthalpy(enthalpyMap.reduced_size());

      enthalpy = 0.0;

      auto dependencies =
          prepared_hydrostatic_analytic_solve_test_utils::make_dependencies();

      mean_field::operators::PreparedHydrostaticEquilibriumOperator
          preparedOperator(f, *f.domainMapperStateless);

      const auto initialReport = preparedOperator.Prepare(
          prepared_hydrostatic_analytic_solve_test_utils::make_state(
              enthalpy, gravityPotential, displacement),
          dependencies, rotation);

      REQUIRE(initialReport.preparedResidual);
      REQUIRE(initialReport.preparedAlgebraicJacobianBlocks);

      mfem::Vector initialResidual;

      preparedOperator.BuildResidual(initialResidual);

      const double initialResidualNorm =
          gravity_prepared_test_utils::global_norm(initialResidual,
                                                   communicator);

      REQUIRE(initialResidualNorm > 1.0e-12);

      /*
       * One discrete Newton step:
       *
       *     M_h delta_h = -R_h.
       *
       * The full four-block Bernoulli Jacobian is rectangular
       * and underdetermined in isolation. Freezing Phi, C,
       * rotation, and displacement makes this a well-defined
       * enthalpy solve.
       */
      prepared_hydrostatic_analytic_solve_test_utils::EnthalpyJacobianOperator
          enthalpyJacobian(enthalpyMap.reduced_size(), preparedOperator);

      mfem::Vector rightHandSide(initialResidual);
      rightHandSide *= -1.0;

      mfem::Vector enthalpyCorrection(enthalpyMap.reduced_size());

      enthalpyCorrection = 0.0;

      /*
       * The reduced operator contains only stellar-supported
       * enthalpy DOFs and is positive definite. MINRES remains
       * appropriate for this symmetric system.
       */
      mfem::MINRESSolver linearSolver(communicator);

      linearSolver.SetOperator(enthalpyJacobian);

      linearSolver.SetRelTol(1.0e-13);
      linearSolver.SetAbsTol(1.0e-14);
      linearSolver.SetMaxIter(2000);
      linearSolver.SetPrintLevel(0);

      linearSolver.Mult(rightHandSide, enthalpyCorrection);

      INFO("Linear solver converged = " << linearSolver.GetConverged());

      INFO("Linear solver iterations = " << linearSolver.GetNumIterations());

      INFO("Linear solver final norm = " << linearSolver.GetFinalNorm());

      REQUIRE(linearSolver.GetConverged());

      enthalpy += enthalpyCorrection;

      /*
       * Only the enthalpy state changed. Geometry, rotation,
       * and algebraic Jacobian data must remain reusable.
       */
      ++dependencies.enthalpy.revision;

      const auto solvedReport = preparedOperator.Prepare(
          prepared_hydrostatic_analytic_solve_test_utils::make_state(
              enthalpy, gravityPotential, displacement),
          dependencies, rotation);

      CHECK(solvedReport.contextReport.updatedEnthalpy);

      CHECK(solvedReport.contextReport.preparedBaseState);

      CHECK_FALSE(solvedReport.contextReport.preparedGeometryState);

      CHECK_FALSE(solvedReport.preparedAlgebraicJacobianBlocks);

      mfem::Vector solvedResidual;

      preparedOperator.BuildResidual(solvedResidual);

      const double solvedResidualNorm =
          gravity_prepared_test_utils::global_norm(solvedResidual,
                                                   communicator);

      const double residualReduction = solvedResidualNorm / initialResidualNorm;

      /*
       * Compare the solved field with the continuum analytic
       * enthalpy over stellar elements only.
       *
       * All three mappings have determinant one, so this
       * normalized L2 error is also unchanged by the physical
       * volume transformation.
       */
      mfem::ParGridFunction solvedEnthalpyField(f.enthalpyFes.get());

      mfem::Vector enthalpyTrue(enthalpyMap.full_size());
      enthalpyMap.scatter(enthalpy, enthalpyTrue);
      solvedEnthalpyField.SetFromTrueDofs(enthalpyTrue);

      const double solvedAnalyticError = solvedEnthalpyField.ComputeL2Error(
          exactEnthalpyCoefficient, nullptr, &stellarElementMarker);

      const double relativeSolvedAnalyticError =
          solvedAnalyticError / exactEnthalpyNorm;

      INFO("Deformation determinant = " << deformationDeterminant);

      INFO("Initial weak residual norm = " << initialResidualNorm);

      INFO("Solved weak residual norm = " << solvedResidualNorm);

      INFO("Weak residual reduction = " << residualReduction);

      INFO("Relative analytic projection floor = " << relativeProjectionError);

      INFO("Relative solved analytic L2 error = "
           << relativeSolvedAnalyticError);

      /*
       * The discrete Bernoulli equation must be solved essentially
       * to the linear-solver floor.
       */
      CHECK(residualReduction < 1.0e-10);

      /*
       * The directly projected analytic enthalpy provides a lower
       * representation bound, but it is not the expected solution
       * of the cross-space discrete Bernoulli equation. The latter
       * also contains potential-projection and mapped-space
       * compatibility errors.
       */
      CHECK(relativeSolvedAnalyticError <
            std::max(5.0 * relativeProjectionError, 1.25e-4));

      /*
       * Record that the analytic error remains within one order of
       * magnitude of the direct enthalpy projection floor.
       */
      CHECK(relativeSolvedAnalyticError / relativeProjectionError < 5.0);
    }
  }
}