#include <catch2/catch_test_macros.hpp>

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <map>
#include <string>

#include <mfem.hpp>
#include <mpi.h>

import experiment;
import experiment.stellar_null_space;
import mean_field;
import test_helpers;

namespace {
    [[nodiscard]] double relative_difference(
        const mfem::Vector &computed,
        const mfem::Vector &reference,
        const MPI_Comm communicator
    ) {
        mfem::Vector difference(computed);
        difference -= reference;
        const double scale = std::max(
            {experiment::null_space::global_norm(computed, communicator),
             experiment::null_space::global_norm(reference, communicator), std::numeric_limits<double>::epsilon()}
        );
        return experiment::null_space::global_norm(difference, communicator) / scale;
    }

    void add_block_metrics(
        std::map<
            std::string,
            double> &metrics,
        const std::string &prefix,
        const std::array<
            double,
            6> &norms
    ) {
        for (std::size_t block = 0; block < norms.size(); ++block) {
            metrics.emplace(prefix + experiment::null_space::residualBlockNames[block] + "_norm", norms[block]);
        }
    }
} // namespace

TEST_CASE(
    "Rigid Motion Responses Of The Stellar Equilibrium Jacobian",
    "[null_space][rigid_motion]"
) {
    mean_field::utils::Args args = test_utils::setup_args();
    args.p.rtol                  = 1.0e-12;
    args.p.atol                  = std::min(args.p.atol, 1.0e-14);
    args.p.max_iters             = std::max(args.p.max_iters, 2000);

    experiment::null_space::N3Equilibrium fixture(std::move(args));
    const MPI_Comm communicator = fixture.fem().mesh->GetComm();
    int rank                    = 0;
    MPI_Comm_rank(communicator, &rank);

    const auto modes = experiment::null_space::make_rigid_modes(fixture);
    constexpr std::array<double, 2> rotationFractions{0.0, 0.5};
    constexpr std::array<double, 2> finiteDifferenceSteps{1.0e-4, 1.0e-6};
    const int totalCases = static_cast<int>(rotationFractions.size() * modes.size());
    int completedCases   = 0;

    for (const double rotationFraction : rotationFractions) {
        const mean_field::physics::RigidRotation rotation = fixture.rotation(rotationFraction);
        fixture.prepare(fixture.state(), rotation);

        mfem::Vector constrainedResidual;
        fixture.stellar_operator().BuildResidual(constrainedResidual);
        const mfem::Vector unpinnedResidual = fixture.unpinned_residual();

        REQUIRE(constrainedResidual.Size() == unpinnedResidual.Size());
        REQUIRE(std::isfinite(experiment::null_space::global_norm(constrainedResidual, communicator)));
        REQUIRE(std::isfinite(experiment::null_space::global_norm(unpinnedResidual, communicator)));

        for (const experiment::null_space::RigidMode &mode : modes) {
            experiment::null_space::report_progress(
                communicator, "probing " + mode.name + " at rotation fraction " + std::to_string(rotationFraction) +
                                  " (" + std::to_string(completedCases + 1) + "/" + std::to_string(totalCases) + ")"
            );

            fixture.prepare(fixture.state(), rotation);
            const mfem::Vector unpinnedAction = fixture.unpinned_jacobian_action(mode.direction);
            mfem::Vector constrainedAction;
            fixture.stellar_operator().Mult(mode.direction, constrainedAction);

            mfem::Vector centeringContribution(constrainedAction);
            centeringContribution -= unpinnedAction;

            const double inputNorm       = experiment::null_space::global_norm(mode.direction, communicator);
            const double unpinnedNorm    = experiment::null_space::global_norm(unpinnedAction, communicator);
            const double constrainedNorm = experiment::null_space::global_norm(constrainedAction, communicator);

            REQUIRE(inputNorm > 0.0);
            REQUIRE(std::isfinite(unpinnedNorm));
            REQUIRE(std::isfinite(constrainedNorm));

            std::map<std::string, double> metrics{
                {"input_algebraic_norm", inputNorm},
                {"unpinned_action_norm", unpinnedNorm},
                {"unpinned_action_per_input_norm", unpinnedNorm / inputNorm},
                {"constrained_action_norm", constrainedNorm},
                {"constrained_action_per_input_norm", constrainedNorm / inputNorm},
                {"centering_contribution_norm",
                 experiment::null_space::global_norm(centeringContribution, communicator)},
                {"unpinned_base_residual_norm", experiment::null_space::global_norm(unpinnedResidual, communicator)},
                {"constrained_base_residual_norm",
                 experiment::null_space::global_norm(constrainedResidual, communicator)}
            };

            add_block_metrics(
                metrics, "unpinned_",
                experiment::null_space::residual_block_norms(
                    unpinnedAction, fixture.stellar_operator().GetLayout(), communicator
                )
            );
            add_block_metrics(
                metrics, "constrained_",
                experiment::null_space::residual_block_norms(
                    constrainedAction, fixture.stellar_operator().GetLayout(), communicator
                )
            );

            for (const double step : finiteDifferenceSteps) {
                mfem::Vector plusState(fixture.state());
                plusState.Add(step, mode.direction);
                fixture.prepare(plusState, rotation);
                const mfem::Vector plusResidual = fixture.unpinned_residual();

                mfem::Vector minusState(fixture.state());
                minusState.Add(-step, mode.direction);
                fixture.prepare(minusState, rotation);
                const mfem::Vector minusResidual = fixture.unpinned_residual();

                mfem::Vector finiteDifference(plusResidual);
                finiteDifference -= minusResidual;
                finiteDifference /= 2.0 * step;

                const std::string stepName = step == finiteDifferenceSteps.front() ? "1e-4" : "1e-6";
                metrics.emplace(
                    "finite_difference_relative_error_" + stepName,
                    relative_difference(unpinnedAction, finiteDifference, communicator)
                );
            }

            fixture.prepare(fixture.state(), rotation);

            if (rank == 0) {
                experiment::record_experiment_result(
                    "stellar_rigid_motion_null_space", mode.name,
                    {{"mode_kind",
                      mode.kind == experiment::null_space::RigidModeKind::translation ? "translation" : "rotation"},
                     {"axis", std::to_string(mode.axis)},
                     {"rotation_fraction_of_keplerian", std::to_string(rotationFraction)},
                     {"mesh_file", test_utils::setup_args().mesh_file},
                     {"local_state_dofs", std::to_string(fixture.stellar_operator().Width())}},
                    std::move(metrics)
                );
            }

            ++completedCases;
            experiment::null_space::report_progress(
                communicator,
                "completed " + std::to_string(completedCases) + "/" + std::to_string(totalCases) + " rigid-mode cases"
            );
        }
    }

    experiment::null_space::report_progress(communicator, "rigid-motion probe complete; writing CSV output");
}