#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include <cmath>
#include <mfem.hpp>

import mean_field;
import test_helpers;

using namespace mean_field;
using Catch::Matchers::WithinAbs;
namespace prepared_test = gravity_prepared_test_utils;

TEST_CASE("Prepared Mapped Hdiv Mass Matches Stateless Kernel",
          tags::gravity_prepared) {
  auto args = test_utils::setup_args();
  fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

  operators::PreparedMappedHDivMassOperator prepared_operator(
      f, *f.domainMapperStateless);
  REQUIRE(prepared_operator.Width() ==
          prepared_operator.GetFluxMap().reduced_size());
  REQUIRE(prepared_operator.Height() ==
          prepared_operator.GetFluxMap().reduced_size());

  const mfem::Vector gravity_gradient_true =
      prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(),
                                               0.21);
  const mfem::Vector gravity_gradient =
      prepared_operator.GetFluxMap().gather(gravity_gradient_true);
  const MPI_Comm communicator = f.gravityFluxFes->GetComm();

  mfem::Vector identity_action;
  mfem::Vector deformed_action;

  for (const double deformation_scale : {0.0, 1.0}) {
    const mfem::Vector displacement_true =
        prepared_test::make_displacement(f, deformation_scale);
    const mfem::Vector displacement =
        prepared_operator.GetDisplacementMap().gather(displacement_true);

    prepared_operator.Prepare(displacement);

    mfem::Vector prepared_action;

    prepared_operator.Mult(gravity_gradient, prepared_action);
    mfem::Vector reference_action_true;
    operators::kernels::apply_mapped_hdiv_mass(
        f, *f.domainMapperStateless, gravity_gradient_true, displacement_true,
        reference_action_true);
    const mfem::Vector reference_action =
        prepared_operator.GetFluxMap().gather(reference_action_true);

    const double relative_error = prepared_test::relative_error(
        prepared_action, reference_action, communicator);

    INFO("Deformation scale = " << deformation_scale);
    INFO("Prepared action norm = "
         << prepared_test::global_norm(prepared_action, communicator));
    INFO("Reference action norm = "
         << prepared_test::global_norm(reference_action, communicator));
    INFO("Relative prepared-operator error = " << relative_error);

    REQUIRE(prepared_operator.IsPrepared());
    CHECK_THAT(relative_error, WithinAbs(0.0, 2.0e-11));

    if (deformation_scale == 0.0) {
      identity_action = prepared_action;
    } else {
      deformed_action = prepared_action;
    }
  }

  const double geometry_change = prepared_test::relative_error(
      deformed_action, identity_action, communicator);

  INFO("Relative action change under deformation = " << geometry_change);

  CHECK(prepared_operator.GetPreparationCount() == 2);
  CHECK(geometry_change > 1.0e-5);
}

TEST_CASE("Prepared Mapped Hdiv Mass Preserves Operator Identities",
          tags::gravity_prepared) {
  auto args = test_utils::setup_args();
  fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

  operators::PreparedMappedHDivMassOperator prepared_operator(
      f, *f.domainMapperStateless);
  REQUIRE(prepared_operator.Width() ==
          prepared_operator.GetFluxMap().reduced_size());
  REQUIRE(prepared_operator.Height() ==
          prepared_operator.GetFluxMap().reduced_size());
  const mfem::Vector displacement =
      prepared_operator.GetDisplacementMap().gather(
          prepared_test::make_displacement(f, 1.0));
  prepared_operator.Prepare(displacement);

  const mfem::Vector first = prepared_operator.GetFluxMap().gather(
      prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(),
                                               0.17));
  const mfem::Vector second = prepared_operator.GetFluxMap().gather(
      prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(),
                                               0.83));
  const mfem::Vector combination =
      prepared_test::linear_combination(first, 1.7, second, -0.4);

  mfem::Vector first_action;
  mfem::Vector second_action;
  mfem::Vector combination_action;
  mfem::Vector zero_action;

  prepared_operator.Mult(first, first_action);
  prepared_operator.Mult(second, second_action);
  prepared_operator.Mult(combination, combination_action);

  mfem::Vector expected_combination =
      prepared_test::linear_combination(first_action, 1.7, second_action, -0.4);

  mfem::Vector zero(first.Size());
  zero = 0.0;
  prepared_operator.Mult(zero, zero_action);

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

  const double first_second_product =
      prepared_test::global_dot(first, second_action, communicator);
  const double second_first_product =
      prepared_test::global_dot(second, first_action, communicator);
  const double symmetry_error = prepared_test::relative_scalar_error(
      first_second_product, second_first_product);
  const double linearity_error = prepared_test::relative_error(
      combination_action, expected_combination, communicator);
  const double first_energy =
      prepared_test::global_dot(first, first_action, communicator);
  const double second_energy =
      prepared_test::global_dot(second, second_action, communicator);
  const std::uint64_t preparation_count =
      prepared_operator.GetPreparationCount();

  mfem::Vector repeated_action;
  prepared_operator.Mult(first, repeated_action);

  INFO("u^T M v = " << first_second_product);
  INFO("v^T M u = " << second_first_product);
  INFO("Relative symmetry error = " << symmetry_error);
  INFO("Relative linearity error = " << linearity_error);
  INFO("u^T M u = " << first_energy);
  INFO("v^T M v = " << second_energy);

  CHECK_THAT(symmetry_error, WithinAbs(0.0, 2.0e-12));
  CHECK_THAT(linearity_error, WithinAbs(0.0, 2.0e-12));
  CHECK_THAT(prepared_test::global_norm(zero_action, communicator),
             WithinAbs(0.0, 1.0e-14));
  CHECK(first_energy > 0.0);
  CHECK(second_energy > 0.0);
  CHECK(prepared_test::relative_error(repeated_action, first_action,
                                      communicator) < 2.0e-14);
  CHECK(prepared_operator.GetPreparationCount() == preparation_count);
}

TEST_CASE("Prepared Mapped Hdiv Mass Diagonal Is Positive Across Both Domains",
          tags::gravity_prepared) {
  auto args = test_utils::setup_args();
  fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

  operators::PreparedMappedHDivMassOperator prepared_operator(
      f, *f.domainMapperStateless);
  const mfem::Vector displacement =
      prepared_operator.GetDisplacementMap().gather(
          prepared_test::make_displacement(f, 1.0));
  prepared_operator.Prepare(displacement);

  mfem::Vector diagonal;
  mfem::Vector true_diagonal;
  prepared_operator.AssembleDiagonal(diagonal);
  prepared_operator.AssembleTrueDiagonal(true_diagonal);

  REQUIRE(diagonal.Size() == prepared_operator.Height());
  REQUIRE(true_diagonal.Size() == prepared_operator.GetFluxMap().full_size());

  const mfem::Vector gathered_true_diagonal =
      prepared_operator.GetFluxMap().gather(true_diagonal);

  for (int i = 0; i < diagonal.Size(); ++i) {
    REQUIRE(std::isfinite(diagonal(i)));
    CHECK(diagonal(i) > 0.0);
    CHECK_THAT(diagonal(i), WithinAbs(gathered_true_diagonal(i),
                                      1.0e-14 * std::abs(diagonal(i))));
  }
}