#include <chrono>
#include <exception>
#include <iostream>
#include <numbers>
#include <stdexcept>
#include <utility>

#include <mfem.hpp>

import mean_field;

using namespace mean_field;

namespace {
    [[nodiscard]] const char *trialDispositionName(
        const solver::nonlinear::LineSearchTrialDisposition disposition
    ) noexcept {
        using Disposition = solver::nonlinear::LineSearchTrialDisposition;
        switch (disposition) {
        case Disposition::accepted:
            return "accepted";
        case Disposition::inadmissible_state:
            return "inadmissible state";
        case Disposition::non_finite_state:
            return "non-finite state";
        case Disposition::non_finite_residual:
            return "non-finite residual";
        case Disposition::insufficient_decrease:
        default:
            return "insufficient decrease";
        }
    }

    [[nodiscard]] const char *linearStatusName(const solver::LinearSolveStatus status) noexcept {
        switch (status) {
        case solver::LinearSolveStatus::converged:
            return "converged";
        case solver::LinearSolveStatus::maximum_iterations:
            return "maximum iterations";
        case solver::LinearSolveStatus::breakdown:
            return "breakdown";
        case solver::LinearSolveStatus::non_finite:
            return "non-finite";
        case solver::LinearSolveStatus::backend_failure:
        default:
            return "backend failure";
        }
    }
} // namespace

int main(
    int argc,
    char **argv
) {
    mfem::Mpi::Init(argc, argv);
    int exitCode = 0;

    try {
        mfem::Device device("cpu");
        int rank = 0;
        MPI_Comm_rank(MPI_COMM_WORLD, &rank);

        utils::Args arguments;
        arguments.mesh_file = "sandbox.smesh";
        arguments.p.rtol    = 1.0e-12;
        arguments.p.atol    = 1.0e-12;
        const auto discretizationStart = std::chrono::steady_clock::now();
        auto finiteElements            = fem::setup_fem(arguments.mesh_file, arguments, 0);
        if (!finiteElements.okay()) {
            throw std::runtime_error("The sandbox could not construct its finite-element discretization.");
        }
        if (rank == 0) {
            std::cout << "Finite-element setup: "
                      << std::chrono::duration<double>(std::chrono::steady_clock::now() - discretizationStart).count()
                      << " s\n";
        }

        constexpr double radius                = utils::RADIUS;
        constexpr double mass                  = utils::MASS;
        constexpr double angularMomentum       = 0.05;
        constexpr double gravitationalConstant = utils::G;
        const double polytropicConstant =
            2.0 * gravitationalConstant * radius * radius / std::numbers::pi_v<double>;
        const double centralDensity =
            std::numbers::pi_v<double> * mass / (4.0 * radius * radius * radius);

        auto stellarModel = model::StellarModel(
            eos::Polytrope({.n = 1.0, .K = polytropicConstant}),
            surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
            integral::FixedTotalMass({.Mtotal = dimensions::MassValue{mass}}),
            integral::FixedAngularMomentum({
                .Jtotal = dimensions::AngularMomentumValue{angularMomentum},
                .axis   = {0.0, 0.0, 1.0},
                .center = {0.0, 0.0, 0.0}
            }),
            constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{centralDensity}})
        );

        auto discretization = equilibrium::makeStellarDiscretization(
            std::move(finiteElements),
            normalization::PhysicalRieszDiagonal{
                dimensions::LengthValue{radius},
                gravitationalConstant
            }
        );

        auto preconditioner = preconditioning::makePreconditioner();
        auto linearSolver   = solver::linear::FGMRES({.restartLength = 40, .printLevel = -1});
        const auto contextStart = std::chrono::steady_clock::now();
        auto context             = solver::makeContext(
            std::move(stellarModel), std::move(discretization), std::move(preconditioner), std::move(linearSolver)
        );
        if (rank == 0) {
            std::cout << "Solver context setup: "
                      << std::chrono::duration<double>(std::chrono::steady_clock::now() - contextStart).count()
                      << " s\n";
        }

        auto observer = solver::nonlinear::makeObserver(
            [](const solver::nonlinear::BeforeIteration &event) {
                int rank = 0;
                MPI_Comm_rank(event.communicator, &rank);
                if (rank == 0) {
                    std::cout << "Newton " << event.iteration << ": |F| = " << event.residualNorm << '\n';
                }
            },
            [](const solver::nonlinear::AfterLineSearchTrial &event) {
                int rank = 0;
                MPI_Comm_rank(event.communicator, &rank);
                if (rank != 0) {
                    return;
                }

                std::cout << "  trial " << event.trial + 1 << ": step = " << event.stepLength
                          << ", outcome = " << trialDispositionName(event.disposition);
                if (!event.rejectionSource.empty()) {
                    std::cout << ", source = " << event.rejectionSource;
                }
                if (event.metric.has_value()) {
                    std::cout << ", |F| = " << event.metric->residualNorm;
                }
                if (event.minimumJacobianDeterminant.has_value()) {
                    std::cout << ", min(det J_map) = " << *event.minimumJacobianDeterminant;
                }
                std::cout << ", prepare = " << event.preparationSeconds << " s"
                          << ", metric = " << event.metricSeconds << " s\n";
            },
            [](const solver::nonlinear::AfterIteration &event) {
                int rank = 0;
                MPI_Comm_rank(event.communicator, &rank);
                if (rank == 0) {
                    std::cout << "  step = " << event.acceptedStepLength
                              << ", trials = " << event.lineSearchTrials
                              << ", |F| = " << event.residualNorm
                              << ", iteration = " << event.iterationSeconds << " s"
                              << ", line search = " << event.lineSearchSeconds << " s"
                              << ", trial preparation = " << event.trialPreparationSeconds << " s"
                              << ", accepted refresh = " << event.preconditionerRefreshSeconds << " s\n";
                    if (event.rollbackSeconds > 0.0) {
                        std::cout << "  rollback = " << event.rollbackSeconds << " s\n";
                    }
                    if (event.linearSolve.has_value()) {
                        const auto &linear = *event.linearSolve;
                        std::cout << "  FGMRES: " << linearStatusName(linear.status)
                                  << ", iterations = " << linear.iterations
                                  << ", restarts = " << linear.restarts
                                  << ", initial/|b| = "
                                  << (linear.rightHandSideNorm > 0.0
                                          ? linear.initialResidualNorm / linear.rightHandSideNorm
                                          : 0.0)
                                  << ", true/|b| = " << linear.relativeTrueResidualNorm
                                  << ", J calls = " << linear.operatorApplications
                                  << ", P^-1 calls = " << linear.inversePreconditionerApplications
                                  << ", solve = " << linear.solveSeconds << " s"
                                  << ", J time = " << linear.operatorSeconds << " s"
                                  << ", P^-1 time = " << linear.inversePreconditionerSeconds << " s\n";
                    }
                }
            }
        );

        auto nonlinearSolver = solver::nonlinear::Newton(solver::nonlinear::NewtonOptions{
            .relativeTolerance = 1.0e-8,
            .absoluteTolerance = 0.0,
            .maximumIterations = 30,
            .linearSolve       = {
                      .relativeTolerance = 3.0e-2,
                      .absoluteTolerance = 0.0,
                      .maximumIterations = 200
            },
            .backtracking = {}
        });

        auto equilibriumSolver  = solver::make(context, nonlinearSolver, observer);
        const auto evaluateStart = std::chrono::steady_clock::now();
        auto report              = equilibriumSolver.evaluate();
        if (rank == 0) {
            std::cout << "Evaluation: "
                      << std::chrono::duration<double>(std::chrono::steady_clock::now() - evaluateStart).count()
                      << " s\n";
            const auto &diagnostics = report.diagnostics();
            std::cout << "Totals: linear = " << diagnostics.totalLinearSolveSeconds
                      << " s, line search = " << diagnostics.totalLineSearchSeconds
                      << " s, trial preparation = " << diagnostics.totalTrialPreparationSeconds
                      << " s, accepted refresh = " << diagnostics.totalPreconditionerRefreshSeconds
                      << " s, rollback = " << diagnostics.totalRollbackSeconds
                      << " s, inadmissible trials = " << diagnostics.inadmissibleLineSearchTrials
                      << ", non-finite trials = " << diagnostics.nonFiniteLineSearchTrials
                      << ", insufficient-decrease trials = " << diagnostics.insufficientDecreaseTrials << '\n';
        }

        if (report.converged()) {
            auto structureView = report.structureView();
            std::cout << "Converged with " << structureView.state().size() << " state values.\n";

            // These are the eventual persistence APIs. Both deliberately
            // throw "not implemented" until the checkpoint schema is chosen:
            // auto structure = structureView.capture();
            // equilibrium::serialize(structureView, "structure.checkpoint");
        } else {
            auto checkpointView = report.lastAcceptedCheckpointView();
            std::cerr << "Solve stopped after " << report.completedNonlinearIterations()
                      << " accepted steps: " << report.failure().message << '\n'
                      << "The last checkpoint has " << checkpointView.state().size() << " state values.\n";
            // auto checkpoint = checkpointView.capture();
            // equilibrium::serialize(checkpointView, "failed-step.checkpoint");
            exitCode = 1;
        }
    } catch (const std::exception &error) {
        std::cerr << "sandbox failure: " << error.what() << '\n';
        exitCode = 2;
    }

    mfem::Mpi::Finalize();
    return exitCode;
}