/**
 * author: Zezhi Wang
 * Copyright (C) 2023 abess-team
 * Licensed under the MIT License.
 */

#pragma once

#include <Eigen/Eigen>
#include <pybind11/pybind11.h>
#include <memory>
#include <utility>

#include <functional>
#include <autodiff/forward/dual.hpp>

using autodiff::dual;
using autodiff::dual2nd;
using Eigen::MatrixXd;
using Eigen::VectorXd;
using Eigen::VectorXi;
using VectorXdual = Eigen::Matrix<dual, -1, 1>;
using VectorXdual2nd = Eigen::Matrix<dual2nd, -1, 1>;
using std::function;
using std::pair;

using ConvexSolver = function<pair<double, VectorXd>(
    function<double(VectorXd const &, pybind11::object const &)>,         // loss_fn
    function<pair<double, VectorXd>(const VectorXd &, pybind11::object)>, // value_and_grad
    const VectorXd &,                                                     // complete_para
    const VectorXi &,                                                     // effective_para_index
    pybind11::object const &                                              // data
    )>;

class UniversalModel;
// UniversalData includes everything about the statistic model like loss, constraints and the statistic data like samples, operations of data.
// In abess project, UniversalData will be an instantiation of UniversalData in template class algorithm, other instantiation of UniversalData often is matrix.
// Thus, UniversalData is just like a matrix in algorithm, and its instance is often denoted as 'x'.
// In order to work like matrix, UniversalData need the help of utility function like X_seg, slice.
class UniversalData
{
    // complete_para: the initial para, often `para` for short
    // activate_para: this is a concept of abess algorithm, which is considered to have an impact on the model
    // effective_para: IMPORTANT! this is a concept of class UniversalData, which is used to simulate the selection operation of extern data sets.
    //                 non-effective_para is like being deleted in the extern data sets and can't be used by any function,
    //                 thus active_para is a subset of effective_para, and effective_para is a subset of complete_para.
    //                 out of Class UniversalData, effective_para is invisible.
private:
    UniversalModel *model;
    ConvexSolver convex_solver;
    Eigen::Index sample_size;
    Eigen::Index model_size;                // length of complete_para
    VectorXi effective_para_index;          // `complete_para[effective_para_index[i]]` is `effective_para[i]`
    Eigen::Index effective_size;            // length of effective_para_index
    std::shared_ptr<pybind11::object> data; // statistic data from user
public:
    UniversalData() = default;
    UniversalData(Eigen::Index model_size, Eigen::Index sample_size, pybind11::object &data, UniversalModel *model, ConvexSolver convex_solver);
    UniversalData slice_by_para(const VectorXi &target_para_index); // used in util func X_seg() and slice()

    Eigen::Index rows() const;                        // getter of sample_size
    Eigen::Index cols() const;                        // getter of effective_para
    const VectorXi &get_effective_para_index() const; // getter of effective_para_index, only used for log
    UniversalData slice_by_sample(const VectorXi &target_sample_index);
    double loss(const VectorXd &effective_para); // compute the loss with effective_para
    double loss_and_gradient(const VectorXd &effective_para, Eigen::VectorXd &gradient);
    void gradient_and_hessian(const VectorXd &effective_para, VectorXd &gradient, MatrixXd &hessian);
    double optimize(VectorXd &effective_para);
};

class UniversalModel
{
    friend class UniversalData;

private:
    // size of para will be match for data
    function<double(VectorXd const &para, pybind11::object const &data)> loss;
    function<dual(VectorXdual const &para, pybind11::object const &data)> gradient_autodiff;
    function<dual2nd(VectorXdual2nd const &para, pybind11::object const &data)> hessian_autodiff;
    function<pair<double, VectorXd>(VectorXd const &para, pybind11::object const &data)> gradient_user_defined;
    function<MatrixXd(VectorXd const &para, pybind11::object const &data)> hessian_user_defined;
    function<pybind11::object(pybind11::object const &old_data, VectorXi const &target_sample_index)> slice_by_sample;
    function<void(pybind11::object const *p)> deleter = [](pybind11::object const *p)
    { delete p; };

public:
    // register callback function
    void set_loss_of_model(function<double(VectorXd const &, pybind11::object const &)> const &);
    void set_gradient_autodiff(function<dual(VectorXdual const &, pybind11::object const &)> const &);
    void set_hessian_autodiff(function<dual2nd(VectorXdual2nd const &, pybind11::object const &)> const &);
    void set_gradient_user_defined(function<pair<double, VectorXd>(VectorXd const &, pybind11::object const &)> const &);
    void set_hessian_user_defined(function<MatrixXd(VectorXd const &, pybind11::object const &)> const &);
    void set_slice_by_sample(function<pybind11::object(pybind11::object const &, VectorXi const &)> const &);
    void set_deleter(function<void(pybind11::object const &)> const &);
};
