#include <algorithm>
#include <charconv>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <filesystem>
#include <format>
#include <fstream>
#include <iomanip>
#include <numeric>
#include <print>
#include <string>
#include <string_view>
#include <system_error>
#include <type_traits>

#include <stroid/IO/mesh.h>
#include "mfem/general/communication.hpp"
#include "serif/solver/equilibrium.hpp"
#include "serif/solver/seed/state.hpp"
#include "serif/tests/solver/archive.hpp"
#include "serif/tests/solver/physical_summary.hpp"
#include "serif/tests/solver/virial_output.hpp"

namespace {
    /** @brief Reject partial numerical arguments rather than silently selecting different physics. */
    template <typename Value>
    [[nodiscard]] Value argument(std::string_view text) {
        Value value{};
        const auto parsed = std::from_chars(text.data(), text.data() + text.size(), value);
        serif::utils::errors::serif_ensure<serif::utils::errors::SERiFErrorCode::INVALID_ARGUMENT>(
            parsed.ec == std::errc{} && parsed.ptr == text.data() + text.size(), "Invalid numerical argument '{}'.", text);
        return value;
    }

    /** @brief Keep failed-solve diagnostics valid JSON when a measurement is not finite. */
    [[nodiscard]] std::string number(double value) {
        return std::isfinite(value) ? std::format("{:.17g}", value) : "null";
    }

    /**
     * @brief Keep the selected equations separate from the meanings and ordering of their shared coordinates.
     * @note Bump the version if equation semantics change without changing the registered equation names.
     */
    template <typename Model>
    [[nodiscard]] std::string formulation_identity(const Model& star) {
        const auto description = serif::model::resolve_equilibrium_equations(star);
        std::string identity = "formulation=stellar_equations_v1";
        description.interior_equations.for_each([&]<typename Equation>(const Equation&) {
            identity += ';';
            identity += Equation::name;
        });
        identity += ";shape=";
        identity += std::remove_cvref_t<decltype(description.shape_force)>::name;
        return identity;
    }
}

/**
 * @brief Solve the scalar formulation from a radial seed and measure every archived physical request.
 * @note The historical archive supplies mesh topology and sampling requests, never physical field values.
 * An optional current-layout checkpoint is independently reevaluated at the declared physics with zero Newton updates.
 * Its preparation record must identify the same Kelvin map, mesh, formulation and coefficient layout.
 * Legacy or unmarked coordinates require a fresh solve; no convergence status is imported.
 * Physical n, K and J are explicit inputs; M=G=1 and P_surface=0 match the reference family.
 * Outputs distinguish algebraic convergence from the independent 1e-6 vacuum virial acceptance gate.
 * accepted_state.txt remains an algebraic checkpoint; only physical_accepted_state.txt passes both checks.
 */
int main(int argc, char** argv) {
    using namespace serif::utils::errors;

    if (argc != 9 && argc != 10) {
        std::println(stderr, "Usage: ester_validation REPORT CASE OUTPUT n K J NEWTON_CAP KRYLOV_CAP [RECHECK_STATE]");
        return 2;
    }
#ifdef MFEM_USE_MPI
    mfem::Mpi::Init(argc, argv);
#else
    std::println(stderr, "ESTER validation requires the default Hypre/MPI stellar inverse.");
    return 2;
#endif
    namespace fixture = serif::tests::solver;
    namespace archive = fixture::archive;
    namespace solver = serif::solver;
    namespace model = serif::model;
    namespace normalization = solver::normalization;

    const std::filesystem::path report = argv[1], output = argv[3];
    const std::string name = argv[2];
    const double n = argument<double>(argv[4]), K = argument<double>(argv[5]), J = argument<double>(argv[6]);
    const int updates = argument<int>(argv[7]), krylov = argument<int>(argv[8]);
    const bool recheck = argc == 10;
    serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(!recheck || updates == 0,
        "A saved-state recheck must use zero Newton updates; its convergence is evaluated afresh, never imported.");
    serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(std::isfinite(n) && n >= 1.0 && n < 3.0 &&
        std::isfinite(K) && K > 0.0 && std::isfinite(J) && J >= 0.0 && updates >= 0 && updates <= 30 &&
        krylov > 0 && krylov <= 600 && !std::filesystem::exists(output),
        "Require 1<=n<3, positive K, nonnegative J, bounded iteration limits, and a fresh output directory.");
    serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(std::filesystem::create_directories(output), "Cannot create fresh output directory.");
    std::ofstream progress(output / "preparation.txt");
    progress << std::setprecision(17) << "archive=" << std::quoted(report.string()) << "\nrequest_case=" << std::quoted(name)
        << "\nn=" << n << "\nK=" << K << "\nM=1\nG=1\nJ=" << J << "\nP_surface=0\nnewton_cap=" << updates
        << "\nkrylov_cap=" << krylov << '\n';
    // A last stage without 'complete' is an incomplete run, including an exception or process failure.
    const auto stage = [&](std::string_view value) {
        progress << "stage=" << value << '\n';
        progress.flush();
        serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(progress.good(), "Cannot preserve preparation progress.");
    };
    stage("metadata");
    const auto source = report / "inputs" / name;
    const auto metadata = archive::metadata(source);
    const auto parameter = [&](const char* key) { return std::stod(metadata.at(key)); };
    serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(parameter("central_phase") == 0.0 && parameter("Psurf") == 0.0 &&
        parameter("M") == 1.0 && parameter("G") == 1.0 && parameter("R_reference") == 1.0 &&
        parameter("surface_degree") == 4.0 && parameter("rotation_axis_x") == 0.0 && parameter("rotation_axis_y") == 0.0 &&
        parameter("rotation_axis_z") == 1.0 && parameter("rotation_center_x") == 0.0 &&
        parameter("rotation_center_y") == 0.0 && parameter("rotation_center_z") == 0.0,
        "The reference requests must use the documented fixed-M/G vacuum branch and rotation convention.");
    stage("mesh_load");
    auto mesh = stroid::IO::LoadStroidMesh((source / "mesh.smesh").string());
    if (!mesh) {
        stage("failed_mesh_load");
        std::println(stderr, "Cannot load the reference mesh: {}", mesh.error());
        return 2;
    }
    serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(mesh->config.r_star.has_value() && mesh->config.r_infinity.has_value() &&
        mesh->config.r_star.value() == 1.0 && mesh->config.r_infinity.value() == 5.0,
        "The bounded archive study requires mesh reference radii one and five.");
    // Match the existing exterior coordinate's mesh radii. Record this changed map before any state import.
    constexpr std::string_view mapping_identity = "mapping=kelvin_mesh_radii_v1,r_star_ref=1,r_inf_ref=5";
    progress << mapping_identity << '\n';
    stage("mapping_identity");
    if (recheck) {
        const auto preparation_path = std::filesystem::path(argv[9]).parent_path() / "preparation.txt";
        std::ifstream preparation(preparation_path);
        archive::require_preparation_identity(preparation, mapping_identity);
    }
    const double reference_radius = mesh->config.r_star.value();
    const double infinity_radius = mesh->config.r_infinity.value();
    const auto resources = std::make_shared<fixture::StellarResources>(std::move(*mesh), reference_radius, infinity_radius);
    const auto finite_elements = fixture::make_discretization(resources, 4);
    const auto star = fixture::make_model(n, K, 1.0, J, 1.0);
    const std::string formulation = formulation_identity(star);
    progress << formulation << '\n';
    stage("formulation_identity");
    if (recheck) {
        std::ifstream preparation(std::filesystem::path(argv[9]).parent_path() / "preparation.txt");
        archive::require_preparation_identity(preparation, formulation);
    }
    stage("radial_seed");
    const auto radial = solver::seed::make_polytropic_seed(star, {.integration = {.step = 0.001}});
    if (!radial) {
        stage("failed_radial_seed");
        std::println(stderr, "Radial seed failed: {}", radial.error());
        return 2;
    }
    stage("radial_fit");
    const auto fitted = solver::seed::fit_radial_seed(finite_elements, *radial);
    if (!fitted) {
        stage("failed_radial_fit");
        std::println(stderr, "Reference-surface fit failed: {}", fitted.error());
        return 2;
    }
    stage("state_projection");
    auto initial = solver::seed::make_polytropic_state(star, finite_elements, *fitted);
    if (!initial) {
        stage("failed_state_projection");
        std::println(stderr, "Radial projection failed: {}", initial.error());
        return 2;
    }
    stage("surface_basis");
    const auto basis_source = std::filesystem::exists(source / "surface_basis.csv") ? source : report / "inputs" / "n1_J0p05";
    const double basis_error = archive::verify_surface_basis(basis_source, finite_elements, *initial);
    const std::string coordinates = archive::coordinate_identity(initial->layout(), archive::mesh_identity(source / "mesh.smesh"));
    progress << coordinates << '\n';
    stage("coordinate_identity");
    if (recheck) {
        stage("saved_state_recheck");
        std::ifstream preparation(std::filesystem::path(argv[9]).parent_path() / "preparation.txt");
        archive::require_preparation_identity(preparation, coordinates);
        const auto values = archive::read_coefficients(argv[9], initial->values().size());
        std::copy(values.begin(), values.end(), initial->values().begin());
        progress << "recheck_state=" << std::quoted(std::filesystem::absolute(argv[9]).string()) << '\n';
        progress.flush();
        serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(progress.good(), "Cannot record the saved-state input.");
    }
    archive::save_coefficients(output / "initial_state.txt", initial->values());
    std::ofstream seed_summary(output / "seed.json");
    seed_summary << "{\"radial_radius\":" << number(radial->scales.radius.value())
        << ",\"radial_mass\":" << number(radial->scales.mass.value())
        << ",\"radial_central_density\":" << number(radial->scales.central_density.value())
        << ",\"fitted_radius\":" << number(fitted->radius.value()) << ",\"fitted_mass\":" << number(fitted->mass.value()) << "}\n";
    seed_summary.close();
    serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(seed_summary.good(), "Cannot write radial-seed diagnostics.");

    stage("normalization");
    const auto reference = normalization::prepare_stellar_reference_scales(serif::dimensions::MassValue{1.0},
        serif::dimensions::LengthValue{1.0}, serif::dimensions::QuantityValue<serif::dimensions::GravitationalConstant>{1.0});
    if (!reference) {
        stage("failed_normalization");
        std::println(stderr, "Normalization failed: {}", reference.error());
        return 2;
    }
    auto configuration = model::make_discretization(model::DiscretizationOptions{
        .model = star, .finite_elements = finite_elements,
        .normalization = normalization::NormalizationOptions{.scheme = *reference, .coordinates = normalization::Dimensionless{},
            .gravitational_constant = serif::dimensions::QuantityValue<serif::dimensions::GravitationalConstant>{1.0}}
    });
    const solver::SolverOptions options{
        .method = {.relative_tolerance = 0.0, .absolute_tolerance = 1.0e-10, .maximum_iterations = updates,
            .linear_solve = {.relative_tolerance = 0.003, .absolute_tolerance = 0.0, .maximum_iterations = krylov}},
        .linear_solver = {.restart_length = 160, .print_level = 1}
    };
    const auto observe = [&](const solver::SolveEvent& event, const auto& state, const auto& system) {
        std::println("event {}, iteration {}, trial {}, residual {}, step {}: {}", static_cast<int>(event.kind), event.iteration,
            event.trial, event.metric ? number(event.metric->residual_norm) : "unmeasured", event.step, event.message);
        if (system) {
            std::size_t offset = 0;
            for (const auto size : state.layout().residual_sizes) {
                const auto first = system->residual().begin() + offset;
                std::println("row offset {}, count {}, norm {}", offset, size, std::sqrt(std::inner_product(first, first + size, first, 0.0)));
                offset += size;
            }
        }
        if (event.kind == solver::EventKind::accepted) {
            archive::save_coefficients(output / "progress_state.txt", state.values());
        }
        std::fflush(stdout);
    };
    stage("nonlinear_solve");
    const auto started = std::chrono::steady_clock::now();
    const auto result = solver::solve_equilibrium(std::move(configuration), std::move(*initial), options, observe);
    const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - started).count();
    const auto& diagnostics = result.diagnostics();
    stage("vacuum_virial");
    std::expected<fixture::VacuumVirialAssessment, std::string> virial = std::unexpected("The nonlinear solve returned no final state.");
    if (result.has_state()) {
        archive::save_coefficients(output / "last_state.txt", result.state().values());
        if (result.converged()) {
            archive::save_coefficients(output / "accepted_state.txt", result.state().values());
        }
        // The driver enforces the physical z axis and origin above; M=R_reference=G=1 fixes the units.
        const double omega = result.state().template get<fixture::Omega>()[0];
        const serif::equations::RigidRotation rotation{{0.0, 0.0, omega}};
        virial = fixture::vacuum_virial(finite_elements, result.state(), star.equation_of_state(), rotation);
    }
    const bool physical_accepted = result.converged() && virial && virial->passes();
    fixture::write_vacuum_virial(output, virial, result.converged());
    std::ofstream summary(output / "solve.json");
    summary << std::setprecision(17) << "{\"n\":" << n << ",\"K\":" << K << ",\"M\":1,\"G\":1,\"J\":" << J
        << ",\"P_surface\":0,\"formulation\":\"density_primal\",\"initial_kind\":\""
        << (recheck ? "saved_state_recheck" : "fitted_radial_seed") << "\",\"sample_units\":\"canonical\""
        << ",\"formulation_identity\":" << std::quoted(formulation)
        << ",\"seed_case\":\"" << name << "\""
        << ",\"converged\":" << (result.converged() ? "true" : "false") << ",\"status\":" << static_cast<int>(result.status())
        << ",\"converged_meaning\":\"algebraic\",\"accepted_state_meaning\":\"algebraic_checkpoint\""
        << ",\"physical_accepted\":" << (physical_accepted ? "true" : "false")
        << ",\"vacuum_virial_valid\":" << (virial ? "true" : "false")
        << ",\"energy_virial_signed\":" << (virial ? number(virial->samples.back().energy_virial()) : "null")
        << ",\"force_virial_signed\":" << (virial ? number(virial->samples.back().force_virial()) : "null")
        << ",\"virial_quadrature_allowance\":" << (virial ? number(virial->quadrature_allowance) : "null")
        << ",\"virial_tolerance\":1e-6,\"virial_integration_allowance_limit\":1e-7"
        << ",\"newton_cap\":" << updates << ",\"krylov_cap\":" << krylov
        << ",\"nonlinear_tolerance\":1e-10,\"linear_relative_tolerance\":0.003,\"restart_length\":160"
        << ",\"accepted_iterations\":" << diagnostics.accepted_iterations << ",\"attempted_iterations\":" << diagnostics.attempted_iterations
        << ",\"krylov_iterations\":" << diagnostics.total_linear_iterations << ",\"rejected_trials\":" << diagnostics.rejected_trials
        << ",\"last_step\":" << number(diagnostics.last_accepted_step) << ",\"solve_seconds\":" << number(seconds)
        << ",\"initial_residual\":" << number(diagnostics.initial_residual_norm.value_or(-1.0))
        << ",\"final_residual\":" << number(diagnostics.final_residual_norm.value_or(-1.0))
        << ",\"basis_error\":" << number(basis_error) << ",\"basis_source\":\"" << basis_source.filename().string() << "\"";
    if (diagnostics.last_linear_solve) {
        summary << ",\"last_linear_relative_residual\":" << number(diagnostics.last_linear_solve->relative_true_residual_norm);
    }
    summary << "}\n";
    summary.close();
    serif_ensure<SERiFErrorCode::INVALID_ARGUMENT>(summary.good(), "Cannot write complete solve diagnostics.");
    std::println("Termination: {}", result.message());
    stage("physical_sampling");
    if (result.has_state()) {
        if (physical_accepted) {
            archive::save_coefficients(output / "physical_accepted_state.txt", result.state().values());
        }
        archive::sample(report / "meanfield_runs" / name / "samples.csv", output / "samples.csv", finite_elements, result.state(), n, K);
        archive::sample_summary(output, finite_elements, result.state(), report / "meanfield_runs" / name / "surface.csv");
    }
    if (physical_accepted) { stage("complete_physical_accepted"); }
    else if (result.converged()) { stage("complete_algebraically_converged_physical_rejected"); }
    else { stage("complete_unconverged"); }
    std::println("Physical acceptance: {} (full vacuum virial tolerance 1e-6)", physical_accepted ? "passed" : "rejected");
    return physical_accepted ? 0 : 1;
}