// STL includes
#include <expected>

// Library includes
#include <hccd/PathTracker.hpp>
#include <hccd/util.hpp>
#include <spdlog/spdlog.h>

// Project includes
#include "argparse.hpp"


//
//
// Homotopy definition
//
//


///
/// @brief Helper type implementing necessary functions for PathTracker
/// @details Toy example homotopy:
///   G = [[x1],
///        [x2]]
///
///   F = [[x1 + x2 ],
///        [x2 + 0.5]]
///
///   H = (1-t)*G + t*F
///
/// @details Note that
///   y := [[x1],
///         [x2],
///         [t]]
///
struct ToyHomotopy {
    ///
    /// @brief Evaluate H at y
    ///
    static Eigen::VectorXd evaluate_H(Eigen::VectorXd y) {
        Eigen::VectorXd result(2);

        double x1 = y(0);
        double x2 = y(1);
        double t  = y(2);

        result(0) = x1 + t * x2;
        result(1) = x2 + t * 0.5;

        return result;
    }

    ///
    /// @brief Evaluate Jacobian of H at y
    ///
    static Eigen::MatrixXd evaluate_DH(Eigen::VectorXd y) {
        Eigen::MatrixXd result(2, 3);

        double x1 = y(0);
        double x2 = y(1);
        double t  = y(2);

        result(0, 0) = 1;
        result(0, 1) = t;
        result(0, 2) = x2;
        result(1, 0) = 0;
        result(1, 1) = 1;
        result(1, 2) = 0.5;

        return result;
    }
};


//
//
// Perform path tracking
//
//


void tracker_example(hccd::Settings settings, std::size_t num_iterations,
                     std::string output = "temp.csv") {

    hccd::PathTracker<ToyHomotopy> tracker{settings};

    std::ofstream out;
    if (output != "") {
        out.open(output);
        out << "x1b," << "x2b," << "tb," << "x1p," << "x2p," << "tp," << "x1e,"
            << "x2e," << "te," << "x1n," << "x2n," << "tn," << std::endl;
    }

    Eigen::VectorXd y = Eigen::VectorXd::Zero(3);

    for (int i = 0; i < num_iterations; ++i) {
        spdlog::info("Iteration {}", i);

        auto res = tracker.transparent_step(y);
        if (!res) {
            spdlog::error(
                "Newton corrector failed to converge on iteration {} ", i);
            std::terminate();
        }

        Eigen::VectorXd y_start, y_prime, y_hat_e;
        std::tie(y_start, y_prime, y_hat_e, y) = res.value();

        spdlog::info("y:{}", y);

        if (output != "") {
            out << y_start(0) << "," << y_start(1) << "," << y_start(2) << ","
                << y_prime(0) << "," << y_prime(1) << "," << y_prime(2) << ","
                << y_hat_e(0) << "," << y_hat_e(1) << "," << y_hat_e(2) << ","
                << y(0) << "," << y(1) << "," << y(2) << std::endl;
        }
    }
}


//
//
// User Interface
//
//


int main(int argc, char* argv[]) {
    // Parse command line arguments

    argparse::ArgumentParser program("Homotopy continuation path tracker");

    program.add_argument("--verbose")
        .default_value<bool>(false)
        .implicit_value(true);
    program.add_argument("--euler-step-size")
        .help("Step size for Euler predictor")
        .default_value<double>(0.05)
        .scan<'g', double>();
    program.add_argument("--euler-max-tries")
        .help("Maximum number of tries for Euler predictor")
        .default_value<unsigned>(5)
        .scan<'u', unsigned>();
    program.add_argument("--newton-max-iter")
        .help("Maximum number of iterations for Newton corrector")
        .default_value<unsigned>(5)
        .scan<'u', unsigned>();
    program.add_argument("--newton-convergence-threshold")
        .help("Convergence threshold for Newton corrector")
        .default_value<double>(0.01)
        .scan<'g', double>();
    program.add_argument("-s", "--sigma")
        .help("Direction in which the path is traced")
        .default_value<int>(1)
        .scan<'i', int>();

    program.add_argument("-o", "--output")
        .help("Output csv file")
        .default_value<std::string>("");
    program.add_argument("-n", "--num-iterations")
        .help("Number of iterations of the example program to run")
        .default_value<std::size_t>(20)
        .scan<'u', std::size_t>();

    try {
        program.parse_args(argc, argv);
    } catch (const std::runtime_error& err) {
        spdlog::error("{}", err.what());
        std::terminate();
    }

    // Run program

    if (program["--verbose"] == true) spdlog::set_level(spdlog::level::debug);

    hccd::Settings settings{
        .euler_step_size  = program.get<double>("--euler-step-size"),
        .euler_max_tries  = program.get<unsigned>("--euler-max-tries"),
        .newton_max_iter = program.get<unsigned>("--newton-max-iter"),
        .newton_convergence_threshold =
            program.get<double>("--newton-convergence-threshold"),
        .sigma = program.get<int>("--sigma"),
    };

    tracker_example(settings, program.get<std::size_t>("--num-iterations"),
                    program.get<std::string>("--output"));
}