#include <algorithm>
#include <chrono>
#include <filesystem>
#include <memory>
#include <tuple>
#include <type_traits>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "mfem/general/communication.hpp"
#include "serif/backend/mfem/serialize/configuration.hpp"
#include "serif/solver/equilibrium.hpp"
#include "serif/backend/mfem/solver/seed/state.hpp"
#include "serif/solver/seed/polytrope.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "serif/utils/serialize/archive.hpp"
#include "serif/utils/serialize/solver.hpp"

namespace {
    namespace serialization = serif::utils::serialize;
    namespace backend_serialization = serif::backend::mfem::serialize;
    namespace solver = serif::solver;
    namespace fixture = serif::tests::solver;

    /** @brief Keep independently running tests from sharing an archive and remove it on every exit path. */
    struct TemporaryArchive {
        std::filesystem::path path = std::filesystem::temp_directory_path() /
            ("serif-newton-restart-" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()) + ".serif");
        /** @brief Test failures must not leave temporary scientific files behind. */
        ~TemporaryArchive() {
            std::error_code error;
            std::filesystem::remove(path, error);
        }
    };
}

TEST_CASE("A moving stellar checkpoint survives destruction, file loading and Newton continuation", "[serialize][restart][moving_restart]") {
    if (!::mfem::Mpi::IsInitialized()) {
        ::mfem::Mpi::Init();
    }
    using Model = decltype(fixture::make_model());
    std::weak_ptr<const void> original_resources;
    const auto saved = [&] {
        const Model model = fixture::make_model();
        const auto radial = solver::seed::make_polytropic_seed(model);
        REQUIRE(radial);
        const auto finite_elements = serif::model::make_stellar_discretization(model, {.radius = radial->scales.radius},
            serif::geometry::deformation::DomainDeformation<serif::geometry::deformation::SphericalHarmonicSurface>{
                .surface = {.maximum_degree = 0}, .stellar_interior = {}, .vacuum = {}});
        original_resources = finite_elements.owner;
        const auto initial = solver::seed::backend::mfem::make_polytropic_state(model, finite_elements, *radial);
        REQUIRE(initial);
        const auto scales = solver::normalization::prepare_stellar_reference_scales(
            serif::dimensions::MassValue{1.0}, serif::dimensions::LengthValue{1.0});
        REQUIRE(scales);
        const solver::normalization::NormalizationOptions normalization{.scheme = *scales};
        const auto configuration = serif::model::make_discretization(serif::model::DiscretizationOptions{
            .model = model, .finite_elements = finite_elements, .normalization = normalization});
        const solver::SolverOptions options{.method = {.relative_tolerance = 1.0e-8, .absolute_tolerance = 1.0e-10,
            .maximum_iterations = 15, .linear_solve = {.relative_tolerance = 1.0e-3, .maximum_iterations = 200}}};
        const auto uninterrupted = solver::solve_equilibrium(configuration, *initial, options);
        INFO(uninterrupted.message());
        REQUIRE(uninterrupted.converged());
        const auto stop = [](const solver::SolveEvent& event, const auto&, const auto&) {
            return event.kind == solver::EventKind::accepted ? solver::SolveControl::stop : solver::SolveControl::continue_solve;
        };
        const auto interrupted = solver::solve_equilibrium(configuration, *initial, options, stop);
        INFO(interrupted.message());
        REQUIRE(interrupted.status() == solver::SolveStatus::stopped);
        REQUIRE(interrupted.has_checkpoint());
        CHECK(interrupted.diagnostics().accepted_iterations == 1);
        const auto checkpoint = interrupted.checkpoint();
        CHECK(std::ranges::any_of(checkpoint.state.template get<fixture::Shape>(), [](double value) { return value != 0.0; }));
        auto contents = serialization::pack(configuration, options, interrupted);
        return std::tuple{std::move(contents), uninterrupted.state(), uninterrupted.diagnostics(),
            std::type_identity<std::remove_cvref_t<decltype(configuration)>>{}, std::type_identity<std::remove_cvref_t<decltype(options)>>{}};
    }();
    REQUIRE(original_resources.expired());
    using State = std::remove_cvref_t<decltype(std::get<1>(saved))>;
    using Checkpoint = solver::NewtonCheckpoint<State>;
    using Configuration = typename std::remove_cvref_t<decltype(std::get<3>(saved))>::type;
    using Options = typename std::remove_cvref_t<decltype(std::get<4>(saved))>::type;
    const TemporaryArchive archive;
    const auto written = serialization::write(archive.path, std::get<0>(saved), {.author = "SERiF regression", .comment = "Accepted moving-star Newton boundary"});
    REQUIRE(written);
    auto reader = serialization::Reader::open(archive.path);
    REQUIRE(reader);
    const auto checkpoint = reader->read<Checkpoint>("solution/checkpoint");
    if (!checkpoint) {
        INFO(checkpoint.error().path << ": " << checkpoint.error().message);
    }
    REQUIRE(checkpoint);
    const TemporaryArchive text_archive;
    REQUIRE(serialization::write(text_archive.path, serialization::pack(*checkpoint), {.format = serialization::Format::text}));
    auto text_reader = serialization::Reader::open(text_archive.path);
    REQUIRE(text_reader);
    const auto text_checkpoint = text_reader->read<Checkpoint>();
    REQUIRE(text_checkpoint);
    CHECK(serialization::encode(*text_checkpoint) == serialization::encode(*checkpoint));
    const auto configuration_node = reader->node("configuration");
    REQUIRE(configuration_node);
    const auto configuration = backend_serialization::restore_configuration<Configuration>(*configuration_node);
    if (!configuration) {
        INFO(configuration.error().path << ": " << configuration.error().message);
    }
    REQUIRE(configuration);
    const auto options = reader->read<Options>();
    REQUIRE(options);
    const auto resumed = solver::resume_equilibrium(*configuration, *checkpoint, *options);
    INFO(resumed.message());
    REQUIRE(resumed.converged());
    const State& expected = std::get<1>(saved);
    REQUIRE(resumed.state().values().size() == expected.values().size());
    for (std::size_t index = 0; index < expected.values().size(); ++index) {
        CHECK(resumed.state().values()[index] == Catch::Approx(expected.values()[index]).margin(2.0e-10).epsilon(2.0e-10));
    }
    const solver::SolveDiagnostics& diagnostics = std::get<2>(saved);
    CHECK(resumed.diagnostics().initial_residual_norm == diagnostics.initial_residual_norm);
    CHECK(resumed.diagnostics().convergence_threshold == diagnostics.convergence_threshold);
    CHECK(resumed.diagnostics().accepted_iterations == diagnostics.accepted_iterations);
    CHECK(resumed.diagnostics().line_search_trials == diagnostics.line_search_trials);
    CHECK(resumed.diagnostics().rejected_trials == diagnostics.rejected_trials);
    CHECK(resumed.diagnostics().total_linear_iterations == diagnostics.total_linear_iterations);
    CHECK(reader->contains("solution/model"));
    CHECK(reader->contains("solution/reference_scales"));
    CHECK(reader->contains("solution/coordinates"));
    CHECK(reader->read<Checkpoint>("solution/checkpoint"));
    auto corrupted = serialization::encode(*checkpoint);
    corrupted.children.at("progress").children.at("correction") = serialization::encode(std::vector<double>{});
    CHECK_FALSE(serialization::decode<Checkpoint>(corrupted));
    corrupted = serialization::encode(*checkpoint);
    corrupted.children.at("strategy") = serialization::encode(std::string{"unknown_stateful_strategy"});
    CHECK_FALSE(serialization::decode<Checkpoint>(corrupted));
}