#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 {
    class GravityUnknownJacobian final : public mfem::Operator {
    public:
        explicit GravityUnknownJacobian(
            const mean_field::operators::PreparedStellarEquilibriumOperator &stellarOperator
        )
            : mfem::Operator(
                  stellarOperator.GetLayout().size(experiment::null_space::gravityGradientValue) +
                  stellarOperator.GetLayout().size(experiment::null_space::gravityPotentialValue)
              ),
              m_stellarOperator(stellarOperator),
              m_gravityGradientSize(stellarOperator.GetLayout().size(experiment::null_space::gravityGradientValue)) {
            MFEM_VERIFY(Width() == Height(), "The reduced gravity Jacobian must be square.");
        }

        void Mult(
            const mfem::Vector &gravityDirection,
            mfem::Vector &gravityAction
        ) const override {
            MFEM_VERIFY(gravityDirection.Size() == Width(), "The reduced gravity direction has the wrong size.");

            const mfem::Vector gravityGradientDirection(
                const_cast<mfem::real_t *>(gravityDirection.GetData()), m_gravityGradientSize
            );
            const mfem::Vector gravityPotentialDirection(
                const_cast<mfem::real_t *>(gravityDirection.GetData()) + m_gravityGradientSize,
                Width() - m_gravityGradientSize
            );

            m_stellarOperator.GetGravityOperator().ApplyGravityUnknowns(
                gravityGradientDirection,
                gravityPotentialDirection,
                m_stellarOperator.GetGravityContext().GetGeometryContext(),
                gravityAction
            );
        }

        [[nodiscard]] int gravity_gradient_size() const noexcept {
            return m_gravityGradientSize;
        }

    private:
        const mean_field::operators::PreparedStellarEquilibriumOperator &m_stellarOperator;
        int m_gravityGradientSize;
    };

    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]);
        }
    }

    [[nodiscard]] mfem::Vector gravity_residual_blocks(
        const mfem::Vector &completeAction,
        const mean_field::operators::StellarEquilibriumLayout &layout
    ) {
        const mfem::Vector gradient = experiment::null_space::const_residual_view(
            completeAction, layout, experiment::null_space::gravityGradientResidual
        );
        const mfem::Vector potential = experiment::null_space::const_residual_view(
            completeAction, layout, experiment::null_space::gravityPotentialResidual
        );

        mfem::Vector result(gradient.Size() + potential.Size());
        mfem::Vector(result.GetData(), gradient.Size())                    = gradient;
        mfem::Vector(result.GetData() + gradient.Size(), potential.Size()) = potential;
        return result;
    }

    void assign_gravity_completion(
        mfem::Vector &completeDirection,
        const mean_field::operators::StellarEquilibriumLayout &layout,
        const mfem::Vector &gravityCompletion,
        const int gravityGradientSize
    ) {
        const mfem::Vector gravityGradient(
            const_cast<mfem::real_t *>(gravityCompletion.GetData()), gravityGradientSize
        );
        const mfem::Vector gravityPotential(
            const_cast<mfem::real_t *>(gravityCompletion.GetData()) + gravityGradientSize,
            gravityCompletion.Size() - gravityGradientSize
        );
        experiment::null_space::assign_value_block(
            completeDirection, layout, experiment::null_space::gravityGradientValue, gravityGradient
        );
        experiment::null_space::assign_value_block(
            completeDirection, layout, experiment::null_space::gravityPotentialValue, gravityPotential
        );
    }

    void apply_centering_rows(
        const mean_field::operators::PreparedStellarEquilibriumOperator &stellarOperator,
        const mfem::Vector &direction,
        mfem::Vector &action
    ) {
        const auto &layout = stellarOperator.GetLayout();
        const mfem::Vector displacementDirection =
            experiment::null_space::const_value_view(direction, layout, experiment::null_space::displacementValue);
        mfem::Vector displacementAction =
            experiment::null_space::residual_view(action, layout, experiment::null_space::displacementResidual);
        stellarOperator.GetCenteringConstraintOperator().ApplyJacobianRows(displacementDirection, displacementAction);
    }
} // namespace

TEST_CASE(
    "Gravity-Completed Rigid Motion Responses Of The Stellar Equilibrium Jacobian",
    "[null_space][gravity_completed]"
) {
    mean_field::utils::Args args = test_utils::setup_args();
    args.p.rtol                  = 1.0e-11;
    args.p.atol                  = std::min(args.p.atol, 1.0e-13);
    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};
    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);

        GravityUnknownJacobian gravityUnknownJacobian(fixture.stellar_operator());
        mean_field::operators::ReducedGravityFieldPreconditioner gravityPreconditioner(
            fixture.fem(), fixture.stellar_operator().GetGravityContext().GetGeometryContext()
        );

        mfem::MINRESSolver gravitySolver(communicator);
        gravitySolver.SetOperator(gravityUnknownJacobian);
        gravitySolver.SetPreconditioner(gravityPreconditioner);
        gravitySolver.SetRelTol(1.0e-11);
        gravitySolver.SetAbsTol(1.0e-13);
        gravitySolver.SetMaxIter(2000);
        gravitySolver.SetPrintLevel(1);

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

            const mfem::Vector displacementOnlyAction = fixture.unpinned_jacobian_action(mode.direction);
            mfem::Vector gravityRightHandSide =
                gravity_residual_blocks(displacementOnlyAction, fixture.stellar_operator().GetLayout());
            gravityRightHandSide *= -1.0;

            mfem::Vector gravityCompletion(gravityUnknownJacobian.Width());
            gravityCompletion = 0.0;
            gravitySolver.Mult(gravityRightHandSide, gravityCompletion);

            REQUIRE(gravitySolver.GetConverged());

            mfem::Vector gravitySolveAction;
            gravityUnknownJacobian.Mult(gravityCompletion, gravitySolveAction);
            mfem::Vector gravitySolveResidual(gravitySolveAction);
            gravitySolveResidual -= gravityRightHandSide;

            const double gravityRightHandSideNorm =
                experiment::null_space::global_norm(gravityRightHandSide, communicator);
            const double gravitySolveResidualNorm =
                experiment::null_space::global_norm(gravitySolveResidual, communicator);
            const double gravitySolveRelativeResidual =
                gravitySolveResidualNorm / std::max(gravityRightHandSideNorm, std::numeric_limits<double>::epsilon());

            REQUIRE(std::isfinite(gravitySolveRelativeResidual));

            mfem::Vector completedDirection(mode.direction);
            assign_gravity_completion(
                completedDirection, fixture.stellar_operator().GetLayout(), gravityCompletion,
                gravityUnknownJacobian.gravity_gradient_size()
            );

            const mfem::Vector completedUnpinnedAction = fixture.unpinned_jacobian_action(completedDirection);
            mfem::Vector completedConstrainedAction(completedUnpinnedAction);
            apply_centering_rows(fixture.stellar_operator(), completedDirection, completedConstrainedAction);

            mfem::Vector centeringContribution(completedConstrainedAction);
            centeringContribution -= completedUnpinnedAction;

            std::map<std::string, double> metrics{
                {"displacement_only_input_norm", experiment::null_space::global_norm(mode.direction, communicator)},
                {"gravity_completion_norm", experiment::null_space::global_norm(gravityCompletion, communicator)},
                {"completed_input_norm", experiment::null_space::global_norm(completedDirection, communicator)},
                {"displacement_only_action_norm",
                 experiment::null_space::global_norm(displacementOnlyAction, communicator)},
                {"gravity_completed_unpinned_action_norm",
                 experiment::null_space::global_norm(completedUnpinnedAction, communicator)},
                {"gravity_completed_constrained_action_norm",
                 experiment::null_space::global_norm(completedConstrainedAction, communicator)},
                {"centering_contribution_norm",
                 experiment::null_space::global_norm(centeringContribution, communicator)},
                {"gravity_solve_rhs_norm", gravityRightHandSideNorm},
                {"gravity_solve_residual_norm", gravitySolveResidualNorm},
                {"gravity_solve_relative_residual", gravitySolveRelativeResidual},
                {"gravity_solve_iterations", static_cast<double>(gravitySolver.GetNumIterations())},
                {"gravity_solve_final_norm", gravitySolver.GetFinalNorm()}
            };

            add_block_metrics(
                metrics, "displacement_only_",
                experiment::null_space::residual_block_norms(
                    displacementOnlyAction, fixture.stellar_operator().GetLayout(), communicator
                )
            );
            add_block_metrics(
                metrics, "gravity_completed_unpinned_",
                experiment::null_space::residual_block_norms(
                    completedUnpinnedAction, fixture.stellar_operator().GetLayout(), communicator
                )
            );
            add_block_metrics(
                metrics, "gravity_completed_constrained_",
                experiment::null_space::residual_block_norms(
                    completedConstrainedAction, fixture.stellar_operator().GetLayout(), communicator
                )
            );

            if (rank == 0) {
                experiment::record_experiment_result(
                    "gravity_completed_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) +
                                  " gravity-completed rigid-mode cases"
            );
        }
    }

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