#pragma once


// STL includes
#include <expected>

// Library includes
#include <Eigen/Dense>
#include <spdlog/spdlog.h>

// Project includes
#include "util.hpp"


namespace hccd {


namespace detail {
template <typename T>
concept homotopy_c = requires(T) {
    { T::evaluate_H(Eigen::VectorXd()) } -> std::same_as<Eigen::VectorXd>;
    { T::evaluate_DH(Eigen::VectorXd()) } -> std::same_as<Eigen::MatrixXd>;
};
} // namespace detail


///
/// @brief Settings for PathTracker
///
struct Settings {
    double   euler_step_size              = 0.05;
    unsigned euler_max_tries              = 5;
    unsigned newton_max_iter              = 5;
    double   newton_convergence_threshold = 0.01;
    int      sigma = 1; ///< Direction in which the path is traced
};

///
/// @brief Path tracker for the homotopy continuation method
/// @details Uses a predictor-corrector scheme to trace a path defined by a
/// homotopy.
/// @details References:
///   [1] T. Chen and T.-Y. Li, “Homotopy continuation method for solving
///       systems of nonlinear and polynomial equations,” Communications in
///       Information and Systems, vol. 15, no. 2, pp. 119–307, 2015
///
/// @tparam homotopy_c Homotopy defining the path
///
template <detail::homotopy_c Homotopy>
class PathTracker {
public:
    enum Error { NewtonNotConverged };

    PathTracker(Settings settings) : m_settings{settings} {
    }

    ///
    /// @brief Perform one predictor-corrector step
    ///
    Eigen::VectorXd step(Eigen::VectorXd y) {
        auto res = transparent_step(y);

        if (res) {
            return res.value().first;
        } else {
            return std::unexpected(res.error());
        }
    }

    ///
    /// @brief Perform one predictor-corrector step, returning intermediate
    /// results
    ///
    std::expected<std::tuple<Eigen::VectorXd, Eigen::VectorXd, Eigen::VectorXd,
                             Eigen::VectorXd>,
                  Error>
    transparent_step(Eigen::VectorXd y) {
        for (int i = 0; i < m_settings.euler_max_tries; ++i) {
            double step_size = m_settings.euler_step_size / (1 << i);

            const auto [y_hat, y_prime] =
                perform_euler_predictor_step(y, step_size);

            auto res = perform_newton_corrector_step(y_hat);

            if (res) return {{y, y_prime, y_hat, res.value()}};
        }

        return std::unexpected(Error::NewtonNotConverged);
    }

private:
    Settings m_settings;

    std::pair<Eigen::VectorXd, Eigen::VectorXd>
    perform_euler_predictor_step(Eigen::VectorXd y, double step_size) {
        /// Obtain y_prime

        Eigen::MatrixXd DH = Homotopy::evaluate_DH(y);
        auto            qr = DH.transpose().colPivHouseholderQr();
        Eigen::MatrixXd Q  = qr.matrixQ();
        Eigen::MatrixXd R  = qr.matrixR();

        Eigen::VectorXd y_prime = Q.col(2);

        spdlog::debug("Q: \t\t{}x{}; det={}", Q.rows(), Q.cols(),
                      Q.determinant());
        spdlog::debug("R.topRows(2): {}x{}; det={}", R.topRows(2).rows(),
                      R.topRows(2).cols(), R.topRows(2).determinant());

        if (sign(Q.determinant() * R.topRows(2).determinant()) !=
            sign(m_settings.sigma))
            y_prime = -y_prime;

        /// Perform prediction

        Eigen::VectorXd y_hat = y + step_size * y_prime;

        return {y_hat, y_prime};
    }

    std::expected<Eigen::VectorXd, Error>
    perform_newton_corrector_step(Eigen::VectorXd y) {
        Eigen::VectorXd prev_y = y;

        for (int i = 0; i < m_settings.newton_max_iter; ++i) {

            /// Perform correction

            Eigen::MatrixXd DH = Homotopy::evaluate_DH(y);
            Eigen::MatrixXd DH_pinv =
                DH.completeOrthogonalDecomposition().pseudoInverse();

            y = y - DH_pinv * Homotopy::evaluate_H(y);

            /// Check stopping criterion

            spdlog::debug("Newton iteration {}: ||y-prev_y||={}", i,
                          (y - prev_y).norm());
            if ((y - prev_y).norm() < m_settings.newton_convergence_threshold)
                return y;

            prev_y = y;
        }

        return std::unexpected(Error::NewtonNotConverged);
    }

    template <typename T>
    static int sign(T val) {
        return -1 * (val < T(0)) + 1 * (val >= T(0));
    }
};


} // namespace hccd