#pragma once

#include <cstring>

#include <vector>
#include <algorithm>
#include <functional>
#include <memory>

#include "c_common.hpp"
#include "cy_array.hpp"
#include "c_events.hpp"
#include "c_brentq.hpp"

// !!!
// Comment the following 
//typedef DiffeqFuncType DiffeqMethod;
// Comment this import if working outside of CyRK and you just want the program to compile and run for testing/developing the C++ only code.
// "pysolver_cyhook_api.h" is generated by Cython when building the CyRK project.
// It is based off of the "pysolver_cyhook.pyx" file. 
// Read more about how C++ can call python functions here:
// https://stackoverflow.com/questions/10126668/can-i-override-a-c-virtual-function-within-python-with-cython
// and here: https://github.com/dashesy/pyavfcam/blob/master/src/avf.pyx#L27
#include <Python.h>
#include "pysolver_cyhook_api.h"

// We need to forward declare the CySolverResult so that the solver can make calls to its methods
class CySolverResult;
class CySolverDense;

struct _object;
typedef _object PyObject;

enum class ODEMethod : int {
    NO_METHOD_SET,
    BASE_METHOD,
    RK_BASE_METHOD,
    RK23,
    RK45,
    DOP853
};

inline const std::map<ODEMethod, std::string> CyrkODEMethods = {
    { ODEMethod::NO_METHOD_SET,
      "An integration method has not been set." },

    { ODEMethod::BASE_METHOD,
      "Integration method set to base class (you should not see this)." },

    { ODEMethod::RK_BASE_METHOD,
      "Integration method set to base class for generalized RK (you should not see this)." },

    { ODEMethod::RK23,
      "Explicit Runge-Kutta method of order 3(2)." },

    { ODEMethod::RK45,
      "Explicit Runge-Kutta method of order 5(4)" },

    { ODEMethod::DOP853,
      "Explicit Runge-Kutta method of order 8. Error is controlled using a combination of 5th and 3rd order interpolators." }
};

struct ProblemConfig {
    // Required configurations
    DiffeqFuncType diffeq_ptr      = nullptr;  // Pointer to the differential equation function
    double t_start                 = 0.0;      // Start time of the integration
    double t_end                   = 0.0;      // End time of the integration
    std::vector<double> y0_vec     = std::vector<double>(PRE_ALLOC_NUMY);  // Initial conditions for the dependent variables

    // Baseline configurations
    std::vector<char> args_vec     = std::vector<char>(0);    // Additional arguments for the solver
    std::vector<double> t_eval_vec = std::vector<double>(0);  // User-provided time evaluation points
    size_t num_extra               = 0;                       // Number of extra variables (if any)
    size_t expected_size           = 512;                     // Expected size of the solution storage
    size_t max_num_steps           = 0;                       // Default maximum number of steps
    size_t max_ram_MB              = 2048;                    // Default maximum RAM in MB
    PreEvalFunc pre_eval_func      = nullptr;                 // Pre-evaluation function for the differential equation
    bool capture_dense_output      = false;                   // Whether to capture dense output
    bool force_retain_solver       = true;                    // Whether to retain the solver after solving

    // Parameters that will be updated during setup.
    bool capture_extra   = false; // Whether to capture extra outputs
    bool t_eval_provided = false; // Whether t_eval was provided by the user
    size_t num_y         = 0;     // Number of dependent variables
    double num_y_dbl     = 0.0;   // Number of dependent variables as a double
    double num_y_sqrt    = 0.0;   // Square root of the number of dependent variables
    size_t num_dy        = 0;     // Total number of dependent variables including extra outputs
    double num_dy_dbl    = 0.0;   // Total number of dependent variables as a double

    // PySolver Specifics
    PyObject* cython_extension_class_instance = nullptr;  // Pointer to Python instance which holds the diffeq.
    DiffeqMethod py_diffeq_method             = nullptr;  // Said python diffeq.

    // Flags
    bool initialized = false;

    // Event data
    bool check_events = false;
    std::vector<Event> events_vec = std::vector<Event>();

    // Solver specific configurations can be added below via overloading the class.

    virtual ~ProblemConfig() {};
    // Constructors
    ProblemConfig();
    ProblemConfig(
        DiffeqFuncType diffeq_ptr_,
        double t_start_,
        double t_end_,
        std::vector<double>& y0_vec_);
    ProblemConfig(
        DiffeqFuncType diffeq_ptr_,
        double t_start_,
        double t_end_,
        std::vector<double>& y0_vec_,
        std::vector<char>& args_vec_,
        std::vector<double>& t_eval_vec_,
        size_t num_extra_,
        size_t expected_size_,
        size_t max_num_steps_,
        size_t max_ram_MB_,
        PreEvalFunc pre_eval_func_,
        bool capture_dense_output_,
        bool force_retain_solver_,
        std::vector<Event>& events_vec_);

    // Helper functions
    void update_properties(
        DiffeqFuncType diffeq_ptr_,
        double t_start_,
        double t_end_,
        std::vector<double>& y0_vec_
    );
    void update_properties(
        DiffeqFuncType diffeq_ptr_,
        double t_start_,
        double t_end_,
        std::vector<double>& y0_vec_,
        std::vector<char>& args_vec_,
        std::vector<double>& t_eval_vec_,
        size_t num_extra_,
        size_t expected_size_,
        size_t max_num_steps_,
        size_t max_ram_MB_,
        PreEvalFunc pre_eval_func_,
        bool capture_dense_output_,
        bool force_retain_solver_,
        std::vector<Event>& events_vec_
    );
    virtual void initialize();
    virtual void update_properties_from_config(ProblemConfig* new_config_ptr);
};

struct NowStatePointers 
{
    double* t_now_ptr;
    double* y_now_ptr;
    double* dy_now_ptr;
    
    NowStatePointers():
        t_now_ptr(nullptr), y_now_ptr(nullptr), dy_now_ptr(nullptr) { }

    NowStatePointers(double* t_now_ptr, double* y_now_ptr, double* dy_now_ptr):
        t_now_ptr(t_now_ptr), y_now_ptr(y_now_ptr), dy_now_ptr(dy_now_ptr) { }
};

class CySolverBase {

// Methods
protected:
    virtual CyrkErrorCodes p_additional_setup() noexcept;
    virtual void p_estimate_error() noexcept;
    virtual void p_step_implementation() noexcept;
    inline void p_cy_diffeq() noexcept;
    virtual void p_calc_first_step_size() noexcept;
    CyrkErrorCodes p_check_events() noexcept;

public:
    CySolverBase();
    virtual ~CySolverBase();
    CySolverBase(CySolverResult* storage_ptr_);

    virtual void set_Q_order(size_t* Q_order_ptr);
    virtual void set_Q_array(double* Q_ptr);
    void clear_python_refs();
    void offload_to_temp() noexcept;
    void load_back_from_temp() noexcept;
    CyrkErrorCodes resize_num_y(size_t num_y_, size_t num_dy_);
    virtual CyrkErrorCodes setup();
    inline bool check_status() const;
    void take_step();
    void solve();
    // Diffeq can either be the C++ class method or the python hook diffeq. By default set to C++ version.
    std::function<void(CySolverBase*)> diffeq;
    NowStatePointers get_now_state();

    // PySolver methods
    CyrkErrorCodes set_cython_extension_instance(
        PyObject* cython_extension_class_instance,
        DiffeqMethod py_diffeq_method);
    void py_diffeq();

// Attributes
protected:
    // Result storage
    CySolverResult* storage_ptr = nullptr;

    // Diffeq
    DiffeqFuncType diffeq_ptr = nullptr;

    // Function to send to diffeq which is called before dy is calculated
    PreEvalFunc pre_eval_func = nullptr;

    // Python hooks
    DiffeqMethod py_diffeq_method = nullptr;
    PyObject* cython_extension_class_instance = nullptr;

    // Dependent variables properties
    size_t sizeof_dbl_Ny  = 0;
    size_t sizeof_dbl_Ndy = 0;
    double num_y_dbl      = 0.0;
    double num_y_sqrt     = 0.0;

    // Time properties
    double t_start       = 0.0;
    double t_end         = 0.0;
    double t_delta       = 0.0;
    double t_delta_abs   = 0.0;
    double direction_inf = 0.0;

    // Additional arguments for the diffeq
    size_t size_of_args = 0;
    char* args_ptr      = nullptr;
    
    // t_eval information
    std::vector<double> t_eval_reverse_vec = std::vector<double>(); // Used if t_eval is provided and doing backwards integration
    double* t_eval_ptr      = nullptr;
    size_t t_eval_index_old = 0;
    size_t len_t_eval       = 0;
    bool t_eval_finished    = false;
    bool use_t_eval         = false;

    // Flags
    bool direction_flag    = false;
    bool setup_called      = false;
    bool error_flag        = false;
    bool capture_extra     = false;
    bool check_events_flag = false;

    // Dependent variable pointers
    double* y0_ptr = nullptr;
    // y arrays are stored in vectors of size num_y
    // To initialize the memory pick a value that is probably larger than the expected number of variables.
    // It can always be resized later.
    std::vector<double> y_holder_vec = std::vector<double>(PRE_ALLOC_NUMY * 4);
    //double* y_old_ptr    = nullptr; // This needs to be public
    //double* y_now_ptr    = nullptr; // This needs to be public
    double* y_tmp_ptr    = nullptr;
    double* y_interp_ptr = nullptr;
    // For dy, both the dy/dt and any extra outputs are stored. So the maximum size is `num_y` + `num_extra`
    std::vector<double> dy_holder_vec = std::vector<double>(PRE_ALLOC_NUMY * 4);
    double* dy_old_ptr = nullptr;
    //double* dy_now_ptr = nullptr; // This needs to be public
    double* dy_tmp_ptr  = nullptr;
    double* dy_tmp2_ptr = nullptr;

    // Event data holder
    OptimizeInfo root_finder_data      = OptimizeInfo();
    std::vector<double> event_data_vec = std::vector<double>();  // Equivalent to SciPy's "g" array.
    double* event_checks_old_ptr       = nullptr;
    double termination_root            = 0.0;

    // Integration step information
    size_t max_num_steps = 0;

    // Current state properties
    double t_tmp = 0.0;
    size_t len_t = 0;

public:
    bool use_dense_output            = false;
    bool user_provided_max_num_steps = false;
    bool use_pysolver                = false;

    ODEMethod integration_method = ODEMethod::NO_METHOD_SET;

    // State properties that need to be public mostly so that
    // the dense output object can access them.
    size_t num_y       = 0;
    size_t num_extra   = 0;
    size_t num_dy      = 0;
    size_t num_events  = 0;
    double t_old       = 0.0;
    double t_now       = 0.0;
    double* y_old_ptr  = nullptr;
    double* y_now_ptr  = nullptr;
    double* dy_now_ptr = nullptr;
    std::vector<size_t> active_event_indices_vec = std::vector<size_t>();
};
