#include <algorithm>
#include <catch2/catch_test_macros.hpp>
#include <cmath>
#include <limits>
#include <mfem.hpp>
#include <vector>

import mean_field;
import test_helpers;

namespace {
    bool vector_is_finite(const mfem::Vector &vector) {
        for (int index = 0; index < vector.Size(); ++index) {
            if (!std::isfinite(vector(index))) {
                return false;
            }
        }
        return true;
    }

    mfem::Vector make_deterministic_vector(
        const int size,
        const double phase
    ) {
        mfem::Vector vector(size);
        for (int index = 0; index < size; ++index) {
            const double coordinate = static_cast<double>(index + 1);
            vector(index)           = std::sin(phase + 0.017 * coordinate) + 0.25 * std::cos(0.031 * coordinate);
        }
        return vector;
    }

    double global_dot(
        const mfem::Vector &left,
        const mfem::Vector &right,
        const MPI_Comm communicator
    ) {
        REQUIRE(left.Size() == right.Size());
        const double local = left * right;
        double global      = 0.0;
        REQUIRE(MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, communicator) == MPI_SUCCESS);
        return global;
    }

    double global_norm(
        const mfem::Vector &vector,
        const MPI_Comm communicator
    ) {
        return std::sqrt(global_dot(vector, vector, communicator));
    }
} // namespace

TEST_CASE(
    "MPI Runtime Preserves World And Split Communicator Membership",
    "[mpi][distributed][unit]"
) {
    int rank = 0;
    int size = 1;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    std::vector<int> ranks(static_cast<std::size_t>(size), -1);
    MPI_Allgather(&rank, 1, MPI_INT, ranks.data(), 1, MPI_INT, MPI_COMM_WORLD);

    CHECK(size >= 2);
    for (int expected = 0; expected < size; ++expected) {
        CHECK(ranks[expected] == expected);
    }

    MPI_Comm parity_communicator = MPI_COMM_NULL;
    MPI_Comm_split(MPI_COMM_WORLD, rank % 2, rank, &parity_communicator);

    int parity_size = 0;
    MPI_Comm_size(parity_communicator, &parity_size);
    const int expected_parity_size = (size + 1 - rank % 2) / 2;
    CHECK(parity_size == expected_parity_size);

    MPI_Comm_free(&parity_communicator);
}

TEST_CASE(
    "MPI FEM Setup Partitions Every Element Exactly Once",
    "[mpi][distributed][mesh][integration]"
) {
    const mean_field::utils::Args args = test_utils::setup_args();
    const mean_field::fem::FEM f       = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    const long long local_elements     = f.mesh->GetNE();
    long long global_elements          = 0;
    long long minimum_elements         = 0;

    MPI_Allreduce(&local_elements, &global_elements, 1, MPI_LONG_LONG, MPI_SUM, f.mesh->GetComm());
    MPI_Allreduce(&local_elements, &minimum_elements, 1, MPI_LONG_LONG, MPI_MIN, f.mesh->GetComm());

    CHECK(global_elements == f.smesh.mesh->GetNE());
    CHECK(minimum_elements > 0);
    CHECK(f.logicalReferenceMesh->GetNE() == f.mesh->GetNE());
}

TEST_CASE(
    "MPI Prepared Gravity Operators Preserve Global Algebraic Identities",
    "[mpi][distributed][gravity][operators][unit]"
) {
    const auto args        = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    using GeometryContext  = mean_field::operators::context::gravity_field::GravityFieldGeometryContext;
    GeometryContext geometry_context(f, *f.domainMapperStateless);

    mfem::Vector displacement_true(f.displacementFes->GetTrueVSize());
    displacement_true               = 0.0;
    const mfem::Vector displacement = geometry_context.GetDisplacementMap().gather(displacement_true);
    geometry_context.PreparePrimal(displacement, {0}, {0});

    const mfem::Operator &mass = geometry_context.GetMassOperator();
    const mfem::Vector first   = make_deterministic_vector(mass.Width(), 0.17);
    const mfem::Vector second  = make_deterministic_vector(mass.Width(), 0.83);
    mfem::Vector combination(first);
    combination *= 1.7;
    combination.Add(-0.4, second);

    mfem::Vector first_action;
    mfem::Vector second_action;
    mfem::Vector combination_action;
    mass.Mult(first, first_action);
    mass.Mult(second, second_action);
    mass.Mult(combination, combination_action);

    mfem::Vector expected_combination(first_action);
    expected_combination *= 1.7;
    expected_combination.Add(-0.4, second_action);
    mfem::Vector linearity_difference(combination_action);
    linearity_difference -= expected_combination;

    const MPI_Comm communicator = f.mesh->GetComm();
    const double symmetry_scale = std::max(
        {std::abs(global_dot(first, second_action, communicator)),
         std::abs(global_dot(second, first_action, communicator)), std::numeric_limits<double>::epsilon()}
    );
    const double symmetry_error =
        std::abs(global_dot(first, second_action, communicator) - global_dot(second, first_action, communicator)) /
        symmetry_scale;
    const double linearity_error =
        global_norm(linearity_difference, communicator) /
        std::max(global_norm(expected_combination, communicator), std::numeric_limits<double>::epsilon());

    CHECK(symmetry_error <= 2.0e-12);
    CHECK(linearity_error <= 2.0e-12);

    const mfem::Operator &divergence           = geometry_context.GetDivergenceOperator();
    const mfem::Operator &transpose_divergence = geometry_context.GetTransposeDivergenceOperator();
    const mfem::Vector flux                    = make_deterministic_vector(divergence.Width(), 0.41);
    const mfem::Vector potential               = make_deterministic_vector(divergence.Height(), 0.67);
    mfem::Vector divergence_action;
    mfem::Vector transpose_action;
    divergence.Mult(flux, divergence_action);
    transpose_divergence.Mult(potential, transpose_action);

    const double forward_product   = global_dot(potential, divergence_action, communicator);
    const double transpose_product = global_dot(flux, transpose_action, communicator);
    const double adjoint_scale =
        std::max({std::abs(forward_product), std::abs(transpose_product), std::numeric_limits<double>::epsilon()});
    const double adjoint_error = std::abs(forward_product - transpose_product) / adjoint_scale;

    CHECK(adjoint_error <= 2.0e-12);
}

TEST_CASE(
    "MPI Coupled Gravity LDU Is Stationary Linear And Does Not Reprepare Geometry",
    "[mpi][distributed][gravity][preconditioning][integration]"
) {
    const auto args        = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    using GeometryContext  = mean_field::operators::context::gravity_field::GravityFieldGeometryContext;
    GeometryContext geometryContext(f, *f.domainMapperStateless);

    mfem::Vector displacementTrue(f.displacementFes->GetTrueVSize());
    displacementTrue                = 0.0;
    const mfem::Vector displacement = geometryContext.GetDisplacementMap().gather(displacementTrue);
    geometryContext.PreparePrimal(displacement, {.value = 1}, {.value = 1});

    namespace backend         = mean_field::preconditioning::backend;
    namespace preconditioning = mean_field::preconditioning;
    const auto block          = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, backend::HypreBoomerAMG{backend::FixedCycles{.cycles = 1}},
        preconditioning::GravityApproximateLDU{}
    );
    auto prepared             = preconditioning::prepare(f, geometryContext, block);

    const mfem::Vector first  = make_deterministic_vector(prepared.Width(), 0.23);
    const mfem::Vector second = make_deterministic_vector(prepared.Width(), 0.79);
    mfem::Vector combined(first);
    combined *= 1.3;
    combined.Add(-0.45, second);

    mfem::Vector firstAction(prepared.Height());
    mfem::Vector secondAction(prepared.Height());
    mfem::Vector combinedAction(prepared.Height());
    mfem::Vector repeatedAction(prepared.Height());
    firstAction                            = 0.0;
    secondAction                           = 0.0;
    combinedAction                         = 0.0;
    repeatedAction                         = 0.0;

    const std::uint64_t massPreparations   = geometryContext.GetMassOperator().GetPreparationCount();
    const std::uint64_t sourcePreparations = geometryContext.GetSourceOperator().GetPreparationCount();
    double *const combinedStorage          = combinedAction.GetData();

    prepared.Mult(first, firstAction);
    prepared.Mult(second, secondAction);
    prepared.Mult(combined, combinedAction);
    prepared.Mult(first, repeatedAction);

    mfem::Vector expectedCombined(firstAction);
    expectedCombined *= 1.3;
    expectedCombined.Add(-0.45, secondAction);

    const MPI_Comm communicator = f.mesh->GetComm();
    mfem::Vector linearityDifference(combinedAction);
    linearityDifference -= expectedCombined;
    mfem::Vector determinismDifference(repeatedAction);
    determinismDifference -= firstAction;
    const double linearityError =
        global_norm(linearityDifference, communicator) /
        std::max(global_norm(expectedCombined, communicator), std::numeric_limits<double>::epsilon());

    CHECK(vector_is_finite(combinedAction));
    CHECK(linearityError <= 5.0e-12);
    CHECK(global_norm(determinismDifference, communicator) <= 5.0e-14);
    CHECK(combinedAction.GetData() == combinedStorage);
    CHECK(geometryContext.GetMassOperator().GetPreparationCount() == massPreparations);
    CHECK(geometryContext.GetSourceOperator().GetPreparationCount() == sourcePreparations);
    CHECK(prepared.GetFactorization().GetStatistics().applications == 4);
}

TEST_CASE(
    "MPI Gravity Analysis And Solve Produce Finite Distributed Fields",
    "[mpi][distributed][gravity][integration]"
) {
    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    *f.displacement        = 0.0;

    using DomainSchema     = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    mfem::Vector attribute_density(f.smesh.mesh->attributes.Max());
    attribute_density = 0.0;
    for (int index = 0; index < f.smesh.mesh->attributes.Size(); ++index) {
        const int attribute = f.smesh.mesh->attributes[index];
        if (DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Stellar>(attribute)) {
            attribute_density(attribute - 1) = 1.0;
        }
    }

    mfem::PWConstCoefficient density_coefficient(attribute_density);
    mfem::ParGridFunction density(f.densityFes.get());
    density.ProjectCoefficient(density_coefficient);
    mean_field::analysis::conserve_mass(f, density, mean_field::utils::MASS);

    const double integrated_mass = mean_field::analysis::domain_integrate_grid_function(
        f, density, mean_field::utils::DOMAINS::STELLAR, mean_field::mapping::COORDINATE_SPACE::PHYSICAL
    );
    f.com = mean_field::analysis::get_com(f, density);
    f.Q   = mean_field::physics::compute_quadrupole_moment_tensor(f, density, f.com);

    const mean_field::physics::GravitySolution solution = mean_field::physics::solve_gravity_field(
        f,
        mean_field::physics::GravitySolveOptions{
            .relativeTolerance = 1.0e-12, .absoluteTolerance = 1.0e-15, .maximumIterations = 1000
        },
        density, *f.displacement
    );

    mfem::Vector flux_true;
    mfem::Vector potential_true;
    solution.gradPhi.GetTrueDofs(flux_true);
    solution.phi.GetTrueDofs(potential_true);

    const int local_finite = vector_is_finite(flux_true) && vector_is_finite(potential_true) ? 1 : 0;
    int globally_finite    = 0;
    MPI_Allreduce(&local_finite, &globally_finite, 1, MPI_INT, MPI_MIN, f.mesh->GetComm());

    const double local_norms[2]{flux_true * flux_true, potential_true * potential_true};
    double global_norms[2]{};
    MPI_Allreduce(local_norms, global_norms, 2, MPI_DOUBLE, MPI_SUM, f.mesh->GetComm());

    CHECK(globally_finite == 1);
    CHECK(std::abs(integrated_mass - mean_field::utils::MASS) <= 1.0e-12 * mean_field::utils::MASS);
    CHECK(global_norms[0] > std::numeric_limits<double>::min());
    CHECK(global_norms[1] > std::numeric_limits<double>::min());
}