#pragma once

#include <algorithm>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <limits>
#include <stdexcept>

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

// Include after importing mean_field, as with the other frozen-state probes.
namespace plateau_diagnostics {
    struct GravityApplicationProbeOptions final {
        int directions{3};
    };

    /* Exercise complete, nontrivial compact vectors: isolated basis entries
     * can miss permutation or normalization mistakes shared by large blocks.
     * The reference is deliberately the full normalized root application.
     * No preparation, accepted-state update, or convergence claim occurs here. */
    template <mean_field::solver::GravitySubsystemRuntime Subsystem>
    void ProbeGravityApplications(
        const Subsystem &subsystem,
        const std::filesystem::path &outputDirectory,
        const GravityApplicationProbeOptions options = {}
    ) {
        using namespace mean_field::solver;
        const auto communicator = subsystem.GetCommunicator();
        int rank                = 0;
        MPI_Comm_rank(communicator, &rank);
        const int localValid = subsystem.IsCurrent() && options.directions > 0;
        int allValid         = 0;
        MPI_Allreduce(&localValid, &allValid, 1, MPI_INT, MPI_MIN, communicator);
        int minimumDirections = 0, maximumDirections = 0;
        MPI_Allreduce(&options.directions, &minimumDirections, 1, MPI_INT, MPI_MIN, communicator);
        MPI_Allreduce(&options.directions, &maximumDirections, 1, MPI_INT, MPI_MAX, communicator);
        if (!allValid || minimumDirections != maximumDirections) {
            throw std::invalid_argument(
                "Gravity application probes require a current subsystem and common positive direction count."
            );
        }
        PreparedGravitySubsystem reference(
            subsystem.GetFullOperator(), subsystem.GetFullPreconditioner(), subsystem.GetPartition(), communicator,
            linear::FGMRES{}, RootRestrictedGravityApplications{}
        );
        const int size            = subsystem.GetPartition().GravityValueSize();
        const long long localSize = size;
        long long globalOffset    = 0;
        MPI_Exscan(&localSize, &globalOffset, 1, MPI_LONG_LONG, MPI_SUM, communicator);
        if (rank == 0)
            globalOffset = 0;
        const auto norm = [communicator](const mfem::Vector &vector) {
            const double local = vector * vector;
            double global      = 0.0;
            MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, communicator);
            return std::sqrt(global);
        };
        const auto elapsed = [communicator](const double start) {
            const double local = MPI_Wtime() - start;
            double global      = 0.0;
            MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_MAX, communicator);
            return global;
        };
        std::ofstream output;
        if (rank == 0) {
            std::filesystem::create_directories(outputDirectory);
            output.open(outputDirectory / "gravity_applications.csv");
            output
                << std::setprecision(17)
                << "direction,compact_operator,compact_inverse,operator_norm,operator_error,operator_relative_error,"
                   "compact_operator_seconds,root_operator_seconds,inverse_norm,inverse_error,inverse_relative_error,"
                   "compact_inverse_seconds,root_inverse_seconds\n";
        }
        int outputReady = rank != 0 || output.good();
        MPI_Allreduce(MPI_IN_PLACE, &outputReady, 1, MPI_INT, MPI_MIN, communicator);
        if (!outputReady)
            throw std::runtime_error("Could not create gravity application diagnostics.");
        mfem::Vector direction(size), compact(size), restricted(size), difference(size);
        for (int trial = 0; trial < options.directions; ++trial) {
            for (int i = 0; i < size; ++i) {
                const double coordinate = static_cast<double>(globalOffset + i + 1);
                direction(i) =
                    std::sin(0.37 * coordinate * (trial + 1)) + 0.4 * std::cos(0.11 * coordinate * (trial + 2));
            }
            double start = MPI_Wtime();
            subsystem.GetOperator().Mult(direction, compact);
            const double compactOperatorSeconds = elapsed(start);
            start                               = MPI_Wtime();
            reference.GetOperator().Mult(direction, restricted);
            const double rootOperatorSeconds = elapsed(start);
            difference                       = compact;
            difference -= restricted;
            const double operatorNorm  = norm(restricted);
            const double operatorError = norm(difference);
            start                      = MPI_Wtime();
            subsystem.GetPreconditioner().Mult(direction, compact);
            const double compactInverseSeconds = elapsed(start);
            start                              = MPI_Wtime();
            reference.GetPreconditioner().Mult(direction, restricted);
            const double rootInverseSeconds = elapsed(start);
            difference                      = compact;
            difference -= restricted;
            const double inverseNorm  = norm(restricted);
            const double inverseError = norm(difference);
            const auto relative       = [](double error, double scale) {
                return scale == 0.0 ? (error == 0.0 ? 0.0 : std::numeric_limits<double>::infinity()) : error / scale;
            };
            if (rank == 0) {
                output << trial << ',' << subsystem.UsesCompactOperator() << ','
                       << subsystem.UsesCompactPreconditioner() << ',' << operatorNorm << ',' << operatorError << ','
                       << relative(operatorError, operatorNorm) << ',' << compactOperatorSeconds << ','
                       << rootOperatorSeconds << ',' << inverseNorm << ',' << inverseError << ','
                       << relative(inverseError, inverseNorm) << ',' << compactInverseSeconds << ','
                       << rootInverseSeconds << '\n';
            }
        }
    }
} // namespace plateau_diagnostics