#include <algorithm>
#include <cmath>
#include <functional>
#include <limits>
#include <numbers>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace rotational_displacement_force_analytic_test_utils {
    [[nodiscard]] double global_dot(
        const mfem::Vector &left,
        const mfem::Vector &right,
        MPI_Comm communicator
    ) {
        REQUIRE(left.Size() == right.Size());

        const double localDot = left * right;
        double globalDot      = 0.0;

        MPI_Allreduce(&localDot, &globalDot, 1, MPI_DOUBLE, MPI_SUM, communicator);

        return globalDot;
    }

    [[nodiscard]] double relative_error(
        const double computed,
        const double expected
    ) {
        return std::abs(computed - expected) / std::max(std::abs(expected), std::numeric_limits<double>::epsilon());
    }

    [[nodiscard]] mfem::Vector project_constant_density(
        const mean_field::fem::FEM &f,
        const double densityValue
    ) {
        mfem::ParGridFunction density(f.densityFes.get());
        mfem::ConstantCoefficient densityCoefficient(densityValue);
        density.ProjectCoefficient(densityCoefficient);

        mfem::Vector densityTrue;
        density.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] mfem::Vector project_vector_function(
        const mean_field::fem::FEM &f,
        const std::function<void(
            const mfem::Vector &,
            mfem::Vector &
        )> &function
    ) {
        mfem::ParGridFunction field(f.displacementFes.get());
        mfem::VectorFunctionCoefficient coefficient(3, function);
        field.ProjectCoefficient(coefficient);

        mfem::Vector trueDofs;
        field.GetTrueDofs(trueDofs);
        return trueDofs;
    }
} // namespace rotational_displacement_force_analytic_test_utils

TEST_CASE(
    "Rigid Rotation Gradient And Hessian Action Match The Analytic "
    "Potential",
    tags::rotation_analytic_unit
) {
    mfem::Vector angularVelocity(3);
    angularVelocity(0) = 0.23;
    angularVelocity(1) = -0.31;
    angularVelocity(2) = 0.67;

    mfem::Vector center(3);
    center(0) = 0.11;
    center(1) = -0.07;
    center(2) = 0.05;

    const mean_field::physics::RigidRotation rotation(angularVelocity, center);

    mfem::Vector position(3);
    position(0) = 0.41;
    position(1) = -0.29;
    position(2) = 0.37;

    mfem::Vector direction(3);
    direction(0) = -0.17;
    direction(1) = 0.23;
    direction(2) = 0.13;

    mfem::Vector gradient;
    mfem::Vector hessianAction;

    rotation.potential_gradient(position, gradient);

    rotation.potential_gradient_directional_derivative(direction, hessianAction);

    const double directionalDerivative = rotation.potential_directional_derivative(position, direction);

    CHECK(
        rotational_displacement_force_analytic_test_utils::relative_error(gradient * direction, directionalDerivative) <
        2.0e-15
    );

    constexpr double step = 1.0e-6;

    mfem::Vector plusPosition(position);
    plusPosition.Add(step, direction);

    mfem::Vector minusPosition(position);
    minusPosition.Add(-step, direction);

    mfem::Vector plusGradient;
    mfem::Vector minusGradient;

    rotation.potential_gradient(plusPosition, plusGradient);
    rotation.potential_gradient(minusPosition, minusGradient);

    plusGradient -= minusGradient;
    plusGradient /= 2.0 * step;

    plusGradient -= hessianAction;

    CHECK(plusGradient.Norml2() < 2.0e-10);
}

TEST_CASE(
    "Rotational Displacement Force Reproduces The Homogeneous Sphere "
    "Rotational Virial",
    tags::rotation_analytic_accuracy
) {
    mean_field::utils::Args args = test_utils::setup_args();

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

    REQUIRE(f.okay());

    const double radius       = mean_field::utils::RADIUS;
    const double mass         = mean_field::utils::MASS;
    const double volume       = 4.0 * std::numbers::pi * radius * radius * radius / 3.0;
    const double densityValue = mass / volume;
    const double angularSpeed = 0.73;

    const mfem::Vector density =
        rotational_displacement_force_analytic_test_utils::project_constant_density(f, densityValue);

    mfem::Vector displacement(f.displacementFes->GetTrueVSize());
    displacement = 0.0;

    mfem::Vector angularVelocity(3);
    mfem::Vector center(3);
    angularVelocity    = 0.0;
    center             = 0.0;
    angularVelocity(2) = angularSpeed;

    const mean_field::physics::RigidRotation rotation(angularVelocity, center);

    mfem::Vector residual;

    mean_field::operators::kernels::apply_rotational_displacement_force_residual(
        f, *f.domainMapperStateless, rotation, density, displacement, residual
    );

    const mfem::Vector dilationDirection = rotational_displacement_force_analytic_test_utils::project_vector_function(
        f, [](const mfem::Vector &position, mfem::Vector &value) { value = position; }
    );

    const double computedWork =
        rotational_displacement_force_analytic_test_utils::global_dot(residual, dilationDirection, f.mesh->GetComm());

    const double expectedWork = -(2.0 / 5.0) * mass * angularSpeed * angularSpeed * radius * radius;

    const double relativeError =
        rotational_displacement_force_analytic_test_utils::relative_error(computedWork, expectedWork);

    INFO("Computed rotational virial work = " << computedWork);
    INFO("Analytic rotational virial work = " << expectedWork);
    INFO("Rotational virial relative error = " << relativeError);

    CHECK(computedWork < 0.0);
    CHECK(relativeError < 1.0e-5);
}

TEST_CASE(
    "Rotational Displacement Force Matches The Analytic Off-Axis "
    "Resultant",
    tags::rotation_analytic_accuracy
) {
    mean_field::utils::Args args = test_utils::setup_args();

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

    REQUIRE(f.okay());

    const double radius       = mean_field::utils::RADIUS;
    const double mass         = mean_field::utils::MASS;
    const double volume       = 4.0 * std::numbers::pi * radius * radius * radius / 3.0;
    const double densityValue = mass / volume;
    const double angularSpeed = 0.61;

    const mfem::Vector density =
        rotational_displacement_force_analytic_test_utils::project_constant_density(f, densityValue);

    mfem::Vector displacement(f.displacementFes->GetTrueVSize());
    displacement = 0.0;

    mfem::Vector angularVelocity(3);
    mfem::Vector center(3);
    angularVelocity    = 0.0;
    center             = 0.0;
    angularVelocity(2) = angularSpeed;
    center(0)          = 0.13;
    center(1)          = -0.08;

    const mean_field::physics::RigidRotation rotation(angularVelocity, center);

    mfem::Vector residual;

    mean_field::operators::kernels::apply_rotational_displacement_force_residual(
        f, *f.domainMapperStateless, rotation, density, displacement, residual
    );

    for (int component = 0; component < 3; ++component) {
        const mfem::Vector translationDirection =
            rotational_displacement_force_analytic_test_utils::project_vector_function(
                f, [component](const mfem::Vector &, mfem::Vector &value) {
                    value.SetSize(3);
                    value            = 0.0;
                    value(component) = 1.0;
                }
            );

        const double computedResultant = rotational_displacement_force_analytic_test_utils::global_dot(
            residual, translationDirection, f.mesh->GetComm()
        );

        const double expectedResultant = component < 2 ? mass * angularSpeed * angularSpeed * center(component) : 0.0;

        INFO("Resultant component = " << component);
        INFO("Computed resultant = " << computedResultant);
        INFO("Expected resultant = " << expectedResultant);

        if (component < 2) {
            CHECK(
                rotational_displacement_force_analytic_test_utils::relative_error(
                    computedResultant, expectedResultant
                ) < 1.0e-5
            );
        } else {
            CHECK(std::abs(computedResultant) < 1.0e-10);
        }
    }
}