/**
 * author: Jin Zhu
 * Copyright (C) 2023 abess-team
 * Licensed under the MIT License.
 */
#pragma once

#include <algorithm>
#include <random>
#include <vector>

#include "Algorithm.h"
#include "Data.h"
#include "utilities.h"

class Metric
{
public:
    bool is_cv;
    int Kfold;
    // Eigen::Matrix<Eigen::VectorXd, Dynamic, 1> cv_initial_model_param;
    // Eigen::Matrix<Eigen::VectorXd, Dynamic, 1> cv_initial_coef0;

    std::vector<Eigen::VectorXi> cv_initial_A;
    std::vector<Eigen::VectorXi> cv_initial_I;

    std::vector<Eigen::VectorXi> train_mask_list;
    std::vector<Eigen::VectorXi> test_mask_list;

    std::vector<UniversalData> train_X_list;
    std::vector<UniversalData> test_X_list;
    std::vector<Eigen::MatrixXd> train_y_list;
    std::vector<Eigen::MatrixXd> test_y_list;
    std::vector<Eigen::VectorXd> train_weight_list;
    std::vector<Eigen::VectorXd> test_weight_list;

    std::vector<FIT_ARG> cv_init_fit_arg;

    // std::vector<std::vector<UniversalData>> group_XTX_list;

    std::function<double(double, int, int, int)> ic_method;

    Metric(std::function<double(double, int, int, int)> ic_method, int Kfold = 5) : ic_method(ic_method)
    {
        this->is_cv = Kfold > 1;
        this->Kfold = Kfold;
        if (is_cv)
        {
            cv_init_fit_arg.resize(Kfold);
            train_X_list.resize(Kfold);
            test_X_list.resize(Kfold);
            train_y_list.resize(Kfold);
            test_y_list.resize(Kfold);
            test_weight_list.resize(Kfold);
            train_weight_list.resize(Kfold);
        }
    };

    void set_cv_init_fit_arg(int beta_size, int M)
    {
        for (int i = 0; i < this->Kfold; i++)
        {
            Eigen::VectorXd beta_init;
            Eigen::VectorXd coef0_init;
            coef_set_zero(beta_size, M, beta_init, coef0_init);
            Eigen::VectorXi A_init;
            Eigen::VectorXd bd_init;

            FIT_ARG fit_arg(0, 0., beta_init, coef0_init, bd_init, A_init);

            cv_init_fit_arg[i] = fit_arg;
        }
    }

    // void set_cv_initial_model_param(int Kfold, int p)
    // {
    //   this->cv_initial_model_param = Eigen::MatrixXd::Zero(p, Kfold);
    // };

    // void set_cv_initial_A(int Kfold, int p)
    // {
    //   vector<Eigen::VectorXi> tmp(Kfold);
    //   this->cv_initial_A = tmp;
    // };

    // void set_cv_initial_coef0(int Kfold, int p)
    // {
    //   vector<double> tmp(Kfold);
    //   for (int i = 0; i < Kfold; i++)
    //     tmp[i] = 0;
    //   this->cv_initial_coef0 = tmp;
    // };

    // void update_cv_initial_model_param(Eigen::VectorXd model_param, int k)
    // {
    //   this->cv_initial_model_param.col(k) = model_param;
    // }

    // void update_cv_initial_A(Eigen::VectorXi A, int k)
    // {
    //   this->cv_initial_A[k] = A;
    // }

    // void update_cv_initial_coef0(double coef0, int k)
    // {
    //   this->cv_initial_coef0[k] = coef0;
    // }

    void set_cv_train_test_mask(Data &data, int n, Eigen::VectorXi &cv_fold_id)
    {
        Eigen::VectorXi index_list = Eigen::VectorXi::LinSpaced(n, 0, n - 1);
        auto rule = [&cv_fold_id](int i, int j) -> bool
        { return cv_fold_id(i) < cv_fold_id(j); };
        std::sort(index_list.data(), index_list.data() + index_list.size(), rule);

        int k = 0, st = 0, ed = 1;
        std::vector<Eigen::VectorXi> group_list((unsigned int)this->Kfold);
        while (k < this->Kfold && ed < n)
        {
            int mask = cv_fold_id(index_list(st));
            while (ed < n && mask == cv_fold_id(index_list(ed)))
                ed++;

            group_list[k] = index_list.segment(st, ed - st);
            st = ed;
            ed++;
            k++;
        }

        for (int k = 0; k < this->Kfold; k++)
        {
            std::sort(group_list[k].data(), group_list[k].data() + group_list[k].size());
        }

        // cv train-test partition:
        std::vector<Eigen::VectorXi> train_mask_list_tmp((unsigned int)this->Kfold);
        std::vector<Eigen::VectorXi> test_mask_list_tmp((unsigned int)this->Kfold);
        for (int k = 0; k < this->Kfold; k++)
        {
            int train_x_size = n - group_list[k].size();
            // get train_mask
            Eigen::VectorXi train_mask(train_x_size);
            int i = 0;
            for (int j = 0; j < this->Kfold; j++)
            {
                if (j != k)
                {
                    for (int s = 0; s < group_list[j].size(); s++)
                    {
                        train_mask(i) = group_list[j](s);
                        i++;
                    }
                }
            }
            std::sort(train_mask.data(), train_mask.data() + train_mask.size());
            train_mask_list_tmp[k] = train_mask;
            test_mask_list_tmp[k] = group_list[k];

            slice(data.x, train_mask, this->train_X_list[k]);
            slice(data.x, group_list[k], this->test_X_list[k]);
            slice(data.y, train_mask, this->train_y_list[k]);
            slice(data.y, group_list[k], this->test_y_list[k]);
            slice(data.weight, train_mask, this->train_weight_list[k]);
            slice(data.weight, group_list[k], this->test_weight_list[k]);
        }
        this->train_mask_list = train_mask_list_tmp;
        this->test_mask_list = test_mask_list_tmp;
    };

    // void cal_cv_group_XTX(Data<Eigen::MatrixXd, Eigen::VectorXd, Eigen::VectorXd> &data)
    // {
    //   int p = data.p;
    //   Eigen::VectorXi index = data.g_index;
    //   Eigen::VectorXi gsize = data.g_size;
    //   int N = data.g_num;

    //   std::vector<std::vector<Eigen::MatrixXd>> group_XTX_list_tmp(this->Kfold);

    //   for (int k = 0; k < this->Kfold; k++)
    //   {
    //     int train_size = this->train_mask_list[k].size();
    //     Eigen::MatrixXd train_x(train_size, p);

    //     for (int i = 0; i < train_size; i++)
    //     {
    //       train_x.row(i) = data.x.row(this->train_mask_list[k](i));
    //     };
    //     group_XTX_list_tmp[k] = group_XTX(train_x, index, gsize, train_size, p, N, 1);
    //   }
    //   this->group_XTX_list = group_XTX_list_tmp;
    // }

    double ic(int train_n, int M, int N, Algorithm *algorithm)
    {
        double loss = 2 * (algorithm->get_train_loss() - algorithm->lambda_level * algorithm->beta.cwiseAbs2().sum());
        return this->ic_method(loss, N, algorithm->get_effective_number(), train_n);
    };

    double loss_function(UniversalData &train_x, Eigen::MatrixXd &train_y, Eigen::VectorXd &train_weight, Eigen::VectorXi &g_index,
                         Eigen::VectorXi &g_size, int train_n, int p, int N, Algorithm *algorithm)
    {
        Eigen::VectorXi A = algorithm->get_A_out();
        Eigen::VectorXd beta = algorithm->get_beta();
        Eigen::VectorXd coef0 = algorithm->get_coef0();

        Eigen::VectorXi A_ind = find_ind(A, g_index, g_size, beta.rows(), N);
        UniversalData X_A = X_seg(train_x, train_n, A_ind, algorithm->model_type);

        Eigen::VectorXd beta_A;
        slice(beta, A_ind, beta_A);

        // Eigen::VectorXd beta_A(A_ind.size());
        // for (int k = 0; k < A_ind.size(); k++)
        // {
        //   beta_A(k) = beta(A_ind(k));
        // }
        return algorithm->loss_function(X_A, train_y, train_weight, beta_A, coef0, A, g_index, g_size, 0.0);
    }

    // to do
    Eigen::VectorXd fit_and_evaluate_in_metric(std::vector<Algorithm *> algorithm_list,
                                               Data &data, FIT_ARG &fit_arg)
    {
        Eigen::VectorXd loss_list(this->Kfold);

        if (!is_cv)
        {
            algorithm_list[0]->update_sparsity_level(fit_arg.support_size);
            algorithm_list[0]->update_lambda_level(fit_arg.lambda);
            algorithm_list[0]->update_beta_init(fit_arg.beta_init);
            algorithm_list[0]->update_bd_init(fit_arg.bd_init);
            algorithm_list[0]->update_coef0_init(fit_arg.coef0_init);
            algorithm_list[0]->update_A_init(fit_arg.A_init, data.g_num);

            algorithm_list[0]->fit(data.x, data.y, data.weight, data.g_index, data.g_size, data.n, data.p, data.g_num);

            if (algorithm_list[0]->get_warm_start())
            {
                fit_arg.beta_init = algorithm_list[0]->get_beta();
                fit_arg.coef0_init = algorithm_list[0]->get_coef0();
                fit_arg.bd_init = algorithm_list[0]->get_bd();
            }

            loss_list(0) = this->ic(data.n, data.M, data.g_num, algorithm_list[0]);
        }
        else
        {
            Eigen::VectorXi g_index = data.g_index;
            Eigen::VectorXi g_size = data.g_size;
            int p = data.p;
            int N = data.g_num;

#pragma omp parallel for
            // parallel
            for (int k = 0; k < this->Kfold; k++)
            {
                // get test_x, test_y
                int test_n = this->test_mask_list[k].size();
                int train_n = this->train_mask_list[k].size();

                // train & test data
                // Eigen::MatrixXd train_x = matrix_slice(data.x, this->train_mask_list[k], 0);
                // Eigen::MatrixXd test_x = matrix_slice(data.x, this->test_mask_list[k], 0);
                // Eigen::VectorXd train_y = vector_slice(data.y, this->train_mask_list[k]);
                // Eigen::VectorXd test_y = vector_slice(data.y, this->test_mask_list[k]);
                // Eigen::VectorXd train_weight = vector_slice(data.weight, this->train_mask_list[k]);
                // Eigen::VectorXd test_weight = vector_slice(data.weight, this->test_mask_list[k]);

                // Eigen::VectorXd beta_init;
                algorithm_list[k]->update_sparsity_level(fit_arg.support_size);
                algorithm_list[k]->update_lambda_level(fit_arg.lambda);

                algorithm_list[k]->update_beta_init(this->cv_init_fit_arg[k].beta_init);
                algorithm_list[k]->update_bd_init(this->cv_init_fit_arg[k].bd_init);
                algorithm_list[k]->update_coef0_init(this->cv_init_fit_arg[k].coef0_init);
                algorithm_list[k]->update_A_init(this->cv_init_fit_arg[k].A_init, N);
                // beta_init = this->cv_initial_model_param.col(k).eval();
                // algorithm->update_beta_init(beta_init);
                // algorithm->update_coef0_init(this->cv_initial_coef0[k]);
                // algorithm->update_A_init(this->cv_initial_A[k], N);
                // algorithm->update_train_mask(this->train_mask_list[k]);
                // ??????????????????????????????????????????????????????????????
                algorithm_list[k]->fit(this->train_X_list[k], this->train_y_list[k], this->train_weight_list[k],
                                       g_index, g_size, train_n, p, N);

                if (algorithm_list[k]->get_warm_start())
                {
                    this->cv_init_fit_arg[k].beta_init = algorithm_list[k]->get_beta();
                    this->cv_init_fit_arg[k].coef0_init = algorithm_list[k]->get_coef0();
                    this->cv_init_fit_arg[k].bd_init = algorithm_list[k]->get_bd();
                    // this->update_cv_initial_model_param(algorithm->get_beta(), k);
                    // this->update_cv_initial_A(algorithm->get_A_out(), k);
                    // this->update_cv_initial_coef0(algorithm->get_coef0(), k);
                }

                loss_list(k) =
                    this->loss_function(this->test_X_list[k], this->test_y_list[k], this->test_weight_list[k], g_index,
                                        g_size, test_n, p, N, algorithm_list[k]);
            }
        }

        return loss_list;
    };
};
