#include "rollshap.h"

CharacterVector dimnames_lm_y(const List& input, const int& n_cols_y) {
  
  if (input.size() > 1) {
    
    return input[1];
    
  } else {
    
    CharacterVector result(n_cols_y);
    
    for (int i = 0; i < n_cols_y; i++) {
      
      result[i] = "y";
      result[i] += i + 1;
      
    }
    
    return result;
    
  }
  
}

arma::uvec any_na_x(const NumericMatrix& x) {
  
  int n_rows_x = x.nrow();
  int n_cols_x = x.ncol();
  arma::uvec result(n_rows_x);
  
  for (int i = 0; i < n_rows_x; i++) {
    
    int any_na = 0;
    
    for (int j = 0; (j < n_cols_x) && (any_na == 0); j++) {
      if (std::isnan(x(i, j))) {
        any_na = 1;
      }
    }
    
    result[i] = any_na;
    
  }
  
  return result;
  
}

List roll_lm_z(const SEXP& x, const NumericVector& y,
               const int& width, const arma::vec& weights,
               const bool& intercept, const int& min_obs,
               const bool& complete_obs, const bool& na_restore,
               const bool& online) {
  
  if (Rf_isMatrix(x)) {
    
    NumericMatrix xx(x);
    int n = weights.size();
    int n_rows_xy = xx.nrow();
    int n_cols_x = xx.ncol() + 1;
    arma::uvec arma_any_na(n_rows_xy);
    arma::vec arma_n_obs(n_rows_xy);
    arma::vec arma_sum_w(n_rows_xy);
    arma::mat arma_mean(n_rows_xy, n_cols_x);
    arma::cube arma_cov(n_cols_x, n_cols_x, n_rows_xy);
    arma::vec arma_rsq(n_rows_xy);
    List result(3);
    
    // check 'x' and 'y' arguments for errors
    rollshap::check_rows_equal(n_rows_xy, y.size(), "x", "y");
    
    // check 'width' argument for errors
    rollshap::check_pos_int(width, "width");
    
    // default 'weights' argument is equal-weighted,
    // otherwise check argument for errors
    rollshap::check_weights(n_rows_xy, width, weights, "'x' (and 'y')");
    bool status = rollshap::check_lambda(weights, n_rows_xy, width, online);
    
    // default 'min_obs' argument is 'width',
    // otherwise check argument for errors
    rollshap::check_pos_int(min_obs, "min_obs");
    
    // cbind x and y variables
    NumericMatrix data(n_rows_xy, n_cols_x);
    std::copy(xx.begin(), xx.end(), data.begin());
    std::copy(y.begin(), y.end(), data.begin() + n_rows_xy * (n_cols_x - 1));
    
    // default 'complete_obs' argument is 'true',
    // otherwise check argument for errors
    if (complete_obs) {
      arma_any_na = any_na_x(data);
    } else {
      
      warning("'complete_obs = FALSE' is not supported");
      arma_any_na = any_na_x(data);
      
    }
    
    // compute rolling crossproducts
    if (status && online) {
      
      roll::RollCrossProdOnlineMatXX roll_cov_online(data, n, n_rows_xy, n_cols_x, width,
                                                     weights, intercept, false, min_obs,
                                                     arma_any_na, na_restore,
                                                     arma_n_obs, arma_sum_w, arma_mean,
                                                     arma_cov);
      parallelFor(0, n_cols_x, roll_cov_online);
      
    } else {
      
      roll::RollCrossProdOfflineMatXX roll_cov_offline(data, n, n_rows_xy, n_cols_x, width,
                                                       weights, intercept, false, min_obs,
                                                       arma_any_na, na_restore,
                                                       arma_n_obs, arma_sum_w, arma_mean,
                                                       arma_cov);
      parallelFor(0, n_rows_xy * n_cols_x * (n_cols_x + 1) / 2, roll_cov_offline);
      
    }
    
    // compute rolling linear models
    if (intercept) {
      
      arma::mat arma_coef(n_rows_xy, n_cols_x);
      arma::mat arma_se(n_rows_xy, n_cols_x);
      roll::RollLmMatInterceptTRUE roll_lm_slices(arma_cov, n, n_rows_xy, n_cols_x, width,
                                                  arma_n_obs, arma_sum_w, arma_mean,
                                                  arma_coef, arma_rsq, arma_se);
      parallelFor(0, n_rows_xy, roll_lm_slices);
      
      result = List::create(Named("coefficients") = arma_coef,
                            Named("r.squared") = arma_rsq,
                            Named("std.error") = arma_se);
      
    } else if (!intercept) {
      
      arma::mat arma_coef(n_rows_xy, n_cols_x - 1);
      arma::mat arma_se(n_rows_xy, n_cols_x - 1);
      roll::RollLmMatInterceptFALSE roll_lm_slices(arma_cov, n, n_rows_xy, n_cols_x, width,
                                                   arma_n_obs, arma_sum_w,
                                                   arma_coef, arma_rsq, arma_se);
      parallelFor(0, n_rows_xy, roll_lm_slices);
      
      // create and return a list
      result = List::create(Named("coefficients") = arma_coef,
                            Named("r.squared") = arma_rsq,
                            Named("std.error") = arma_se);
      
    }
    
    return result;
    
  } else {
    
    NumericVector xx(x);
    int n = weights.size();
    int n_rows_xy = xx.size();
    int n_cols_x = 1 + 1;
    arma::uvec arma_any_na(n_rows_xy);
    arma::vec arma_n_obs(n_rows_xy);
    arma::vec arma_sum_w(n_rows_xy);
    arma::mat arma_mean(n_rows_xy, n_cols_x);
    arma::cube arma_cov(n_cols_x, n_cols_x, n_rows_xy);
    arma::vec arma_rsq(n_rows_xy);
    List result(3);
    
    // check 'x' and 'y' arguments for errors
    rollshap::check_rows_equal(n_rows_xy, y.size(), "x", "y");
    
    // check 'width' argument for errors
    rollshap::check_pos_int(width, "width");
    
    // default 'weights' argument is equal-weighted,
    // otherwise check argument for errors
    rollshap::check_weights(n_rows_xy, width, weights, "'x' (and 'y')");
    bool status = rollshap::check_lambda(weights, n_rows_xy, width, online);
    
    // default 'min_obs' argument is 'width',
    // otherwise check argument for errors
    rollshap::check_pos_int(min_obs, "min_obs");
    
    // cbind x and y variables
    NumericMatrix data(n_rows_xy, n_cols_x);
    std::copy(xx.begin(), xx.end(), data.begin());
    std::copy(y.begin(), y.end(), data.begin() + n_rows_xy * (n_cols_x - 1));
    
    // default 'complete_obs' argument is 'true',
    // otherwise check argument for errors
    if (complete_obs) {
      arma_any_na = any_na_x(data);
    } else {
      
      warning("'complete_obs = FALSE' is not supported");
      arma_any_na = any_na_x(data);
      
    }
    
    // compute rolling crossproducts
    if (status && online) {
      
      roll::RollCrossProdOnlineMatXX roll_cov_online(data, n, n_rows_xy, n_cols_x, width,
                                                     weights, intercept, false, min_obs,
                                                     arma_any_na, na_restore,
                                                     arma_n_obs, arma_sum_w, arma_mean,
                                                     arma_cov);
      parallelFor(0, n_cols_x, roll_cov_online);
      
    } else {
      
      roll::RollCrossProdOfflineMatXX roll_cov_offline(data, n, n_rows_xy, n_cols_x, width,
                                                       weights, intercept, false, min_obs,
                                                       arma_any_na, na_restore,
                                                       arma_n_obs, arma_sum_w, arma_mean,
                                                       arma_cov);
      parallelFor(0, n_rows_xy * n_cols_x * (n_cols_x + 1) / 2, roll_cov_offline);
      
    }
    
    // compute rolling linear models
    if (intercept) {
      
      arma::mat arma_coef(n_rows_xy, n_cols_x);
      arma::mat arma_se(n_rows_xy, n_cols_x);
      roll::RollLmMatInterceptTRUE roll_lm_slices(arma_cov, n, n_rows_xy, n_cols_x, width,
                                                  arma_n_obs, arma_sum_w, arma_mean,
                                                  arma_coef, arma_rsq, arma_se);
      parallelFor(0, n_rows_xy, roll_lm_slices);
      
      result = List::create(Named("coefficients") = arma_coef,
                            Named("r.squared") = arma_rsq,
                            Named("std.error") = arma_se);
      
    } else if (!intercept) {
      
      arma::vec arma_coef(n_rows_xy);
      arma::vec arma_se(n_rows_xy);
      roll::RollLmVecInterceptFALSE roll_lm_slices(arma_cov, n, n_rows_xy, width,
                                                   arma_n_obs, arma_sum_w,
                                                   arma_coef, arma_rsq, arma_se);
      parallelFor(0, n_rows_xy, roll_lm_slices);
      
      // create and return a list
      result = List::create(Named("coefficients") = arma_coef,
                            Named("r.squared") = arma_rsq,
                            Named("std.error") = arma_se);
      
    }
    
    return result;
    
  }
  
}

// long double factorial(int n) {
  
//   long double result = 1;
  
//   for (int i = 2; i <= n; i++) {
//     result *= i;
//   }
  
//   return result;
  
// }

// [[Rcpp::export(.roll_shap)]]
SEXP roll_shap(const SEXP& x, const SEXP& y,
               const int& width, const arma::vec& weights,
               const bool& intercept, const int& min_obs,
               const bool& complete_obs, const bool& na_restore,
               const bool& online) {
  
  if (Rf_isMatrix(x) && Rf_isMatrix(y)) {
    
    NumericMatrix xx(x);
    NumericMatrix yy(y);
    
    int n_rows_xy = xx.nrow();
    int n_cols_x = xx.ncol();
    int n_cols_y = yy.ncol();
    int n_combn = 1 << n_cols_x; // 2 ^ n_cols_x
    int n_combn_half = n_combn / 2;
    arma::mat arma_x = arma::mat(xx.begin(), n_rows_xy, n_cols_x);
    arma::ivec arma_n(n_combn);
    arma::umat arma_ix(n_cols_x, n_combn, arma::fill::zeros);
    arma::mat arma_rsq(n_rows_xy, n_combn, arma::fill::zeros);
    arma::mat arma_rsq_sum(n_rows_xy, n_cols_x, arma::fill::zeros);
    List result_z(3);

    // check 'n_cols' variable for errors
    rollshap::check_bounds_int(n_cols_x, 1, 20, "n_cols");

    // number of binary combinations
    for (int k = 0; k < n_combn; k++) {

      int n = 0;

      // find the binary combination
      for (int j = 0; j < n_cols_x; j++) {
        if (k & (1 << j)) {
          
          arma_ix(j, k) = j + 1;

          n += 1;
          
        }
      }
      
      arma_n[k] = n;

    }
    
    arma::field<arma::uvec> arma_subset(n_combn);

    for (int k = 0; k < n_combn; k++) {
      if (arma_n[k] > 0) {
        arma_subset(k) = arma::find(arma_ix.col(k) != 0);
      }
    }

    arma::vec arma_weights(n_cols_x);

    arma_weights[0] = 1.0 / n_cols_x;

    for (int j = 1; j < n_cols_x; j++) {
      arma_weights[j] = arma_weights[j - 1] * j / (n_cols_x - j);
    }

    arma::field<arma::uvec> arma_ix_pos(n_cols_x);
    arma::field<arma::uvec> arma_ix_neg(n_cols_x);
    arma::field<arma::ivec> arma_ix_n(n_cols_x);

    for (int j = 0; j < n_cols_x; j++) {

      arma_ix_pos(j) = arma::find(arma_ix.row(j) != 0);
      arma_ix_neg(j) = arma::find(arma_ix.row(j) == 0);
      arma_ix_n(j) = arma_n(arma_ix_neg(j));

    }

    // create a list of matrices,
    // otherwise a list of lists
    if (n_cols_y == 1) {
      for (int k = 0; k < n_combn; k++) {
        if (arma_n[k] > 0) {
          
          arma::uvec arma_ix_subset = arma_subset(k);
          arma::mat arma_x_subset = arma_x.cols(arma_ix_subset);
          NumericMatrix x_subset(wrap(arma_x_subset));
          
          result_z = roll_lm_z(x_subset, yy(_, 0), width,
                                weights, intercept, min_obs,
                                complete_obs, na_restore,
                                online);
          
          arma::mat arma_rsq_z = result_z["r.squared"];
          arma_rsq.col(k) = arma_rsq_z;
          
        }

      }
        
      // calculate the exact Shapley value for r-squared
      for (int j = 0; j < n_cols_x; j++) {
        
        arma::uvec arma_ix_pos_j = arma_ix_pos(j);
        arma::uvec arma_ix_neg_j = arma_ix_neg(j);
        arma::ivec arma_ix_n_j = arma_ix_n(j);

        for (int k = 0; k < n_combn_half; k++) {
          
          int s = arma_ix_n_j[k];
          
          arma_rsq_sum.col(j) += arma_weights[s] * 
            (arma_rsq.col(arma_ix_pos_j[k]) - arma_rsq.col(arma_ix_neg_j[k]));

        }
        
      }
      
      // create and return a matrix or xts object for Shapley values
      NumericMatrix rsq(wrap(arma_rsq_sum));
      rollshap::xts_attr(rsq, xx);
      
      return rsq;
      
    } else {

      List result_rsq(n_cols_y);
      
      for (int z = 0; z < n_cols_y; z++) {

        arma_rsq.zeros();
        arma_rsq_sum.zeros();

        for (int k = 0; k < n_combn; k++) {

          if (arma_n[k] > 0) {

            arma::uvec arma_ix_subset = arma_subset(k);
            arma::mat arma_x_subset = arma_x.cols(arma_ix_subset);
            NumericMatrix x_subset(wrap(arma_x_subset));

            result_z = roll_lm_z(x_subset, yy(_, z), width,
                                weights, intercept, min_obs,
                                complete_obs, na_restore,
                                online);

            arma::mat arma_rsq_z = result_z["r.squared"];
            arma_rsq.col(k) = arma_rsq_z;

          }

        }
        
        // calculate the exact Shapley value for r-squared
        for (int j = 0; j < n_cols_x; j++) {

          arma::uvec arma_ix_pos_j = arma_ix_pos(j);
          arma::uvec arma_ix_neg_j = arma_ix_neg(j);
          arma::ivec arma_ix_n_j = arma_ix_n(j);
          
          for (int k = 0; k < n_combn_half; k++) {
            
            int s = arma_ix_n_j[k];

            arma_rsq_sum.col(j) += arma_weights[s] *
              (arma_rsq.col(arma_ix_pos_j[k]) - arma_rsq.col(arma_ix_neg_j[k]));
            
          }
          
        }
        
        // create and return a matrix or xts object for Shapley values
        NumericMatrix rsq(wrap(arma_rsq_sum));
        rollshap::xts_attr(rsq, xx);
        
        result_rsq(z) = rsq;
        
      }
      
      // add names to each list
      List dimnames_y = yy.attr("dimnames");
      result_rsq.attr("names") = dimnames_lm_y(dimnames_y, n_cols_y);
      
      return result_rsq;
      
    }
    
  } else if (Rf_isMatrix(x)) {
    
    NumericMatrix xx(x);
    NumericVector yy(y);
    
    int n_rows_xy = xx.nrow();
    int n_cols_x = xx.ncol();
    int n_combn = 1 << n_cols_x; // 2 ^ n_cols_x
    int n_combn_half = n_combn / 2;
    arma::mat arma_x = arma::mat(xx.begin(), n_rows_xy, n_cols_x);
    arma::ivec arma_n(n_combn);
    arma::umat arma_ix(n_cols_x, n_combn, arma::fill::zeros);
    arma::mat arma_rsq(n_rows_xy, n_combn, arma::fill::zeros);
    arma::mat arma_rsq_sum(n_rows_xy, n_cols_x, arma::fill::zeros);
    List result_z(3);

    // check 'n_cols' variable for errors
    rollshap::check_bounds_int(n_cols_x, 1, 20, "n_cols");
    
    // number of binary combinations
    for (int k = 0; k < n_combn; k++) {
      
      int n = 0;
      
      // find the binary combination
      for (int j = 0; j < n_cols_x; j++) {
        if (k & (1 << j)) {
          
          arma_ix(j, k) = j + 1;

          n += 1;       
          
        }        
      }
      
      arma_n[k] = n;
      
    }

    arma::field<arma::uvec> arma_subset(n_combn);

    for (int k = 0; k < n_combn; k++) {
      if (arma_n[k] > 0) {
        arma_subset(k) = arma::find(arma_ix.col(k) != 0);
      }
    }

    arma::vec arma_weights(n_cols_x);

    arma_weights[0] = 1.0 / n_cols_x;

    for (int j = 1; j < n_cols_x; j++) {
      arma_weights[j] = arma_weights[j - 1] * j / (n_cols_x - j);
    }

    arma::field<arma::uvec> arma_ix_pos(n_cols_x);
    arma::field<arma::uvec> arma_ix_neg(n_cols_x);
    arma::field<arma::ivec> arma_ix_n(n_cols_x);

    for (int j = 0; j < n_cols_x; j++) {

      arma_ix_pos(j) = arma::find(arma_ix.row(j) != 0);
      arma_ix_neg(j) = arma::find(arma_ix.row(j) == 0);
      arma_ix_n(j) = arma_n(arma_ix_neg(j));
      
    }

    // number of binary combinations
    for (int k = 0; k < n_combn; k++) {
      
      if (arma_n[k] > 0) {
        
        arma::uvec arma_ix_subset = arma_subset(k);
        arma::mat arma_x_subset = arma_x.cols(arma_ix_subset);
        NumericMatrix x_subset(wrap(arma_x_subset));
        
        result_z = roll_lm_z(x_subset, yy, width,
                             weights, intercept, min_obs,
                             complete_obs, na_restore,
                             online);
        
        arma::mat arma_rsq_z = result_z["r.squared"];
        arma_rsq.col(k) = arma_rsq_z;
        
      }
      
    }
    
    // calculate the exact Shapley value for r-squared
    for (int j = 0; j < n_cols_x; j++) {
      
      arma::uvec arma_ix_pos_j = arma_ix_pos(j);
      arma::uvec arma_ix_neg_j = arma_ix_neg(j);
      arma::ivec arma_ix_n_j = arma_ix_n(j);
      
      for (int k = 0; k < n_combn_half; k++) {
        
        int s = arma_ix_n_j[k];

        arma_rsq_sum.col(j) += arma_weights[s] *
          (arma_rsq.col(arma_ix_pos_j[k]) - arma_rsq.col(arma_ix_neg_j[k]));
        
      }
      
    }
    
    // create and return a matrix or xts object for Shapley values
    NumericMatrix rsq(wrap(arma_rsq_sum));
    rollshap::xts_attr(rsq, xx);
    
    return rsq;
    
  } else if (Rf_isMatrix(y)) {
    
    NumericVector xx(x);
    NumericMatrix yy(y);
    NumericMatrix xxx(xx.size(), 1, xx.begin()); // consistent with roll's roll_lm
    
    int n_rows_xy = xxx.nrow();
    // int n_cols_x = xxx.ncol();
    int n_cols_y = yy.ncol();
    List result_z(3);

    // // check 'n_cols' variable for errors
    // rollshap::check_bounds_int(n_cols_x, 1, 20, "n_cols");
    
    // create a list of matrices,
    // otherwise a list of lists
    if (n_cols_y == 1) {
      
      result_z = roll_lm_z(xxx, yy(_, 0), width,
                           weights, intercept, min_obs,
                           complete_obs, na_restore,
                           online);
      
      arma::vec arma_rsq_z = result_z["r.squared"];
      
      // create and return a vector object for Shapley values
      NumericVector rsq(wrap(arma_rsq_z));
      rsq.attr("dim") = IntegerVector::create(n_rows_xy, 1);
      List dimnames_x = xx.attr("dimnames");
      rollshap::xts_attr(rsq, yy, dimnames_x);
      
      return rsq;
      
    } else {

      List result_rsq(n_cols_y);
      
      for (int z = 0; z < n_cols_y; z++) {
        
        result_z = roll_lm_z(xxx, yy(_, z), width,
                             weights, intercept, min_obs,
                             complete_obs, na_restore,
                             online);
        
        arma::vec arma_rsq_z = result_z["r.squared"];
        
        // create and return a vector object for Shapley values
        NumericVector rsq(wrap(arma_rsq_z));
        rsq.attr("dim") = IntegerVector::create(n_rows_xy, 1);
        List dimnames_x = xx.attr("dimnames");
        rollshap::xts_attr(rsq, yy, dimnames_x);
        
        result_rsq(z) = rsq;
        
      }
      
      // add names to each list
      List dimnames_y = yy.attr("dimnames");
      result_rsq.attr("names") = dimnames_lm_y(dimnames_y, n_cols_y);
      
      return result_rsq;
      
    }
    
  } else {
    
    NumericVector xx(x);
    NumericVector yy(y);
    
    int n_rows_xy = xx.size();
    // int n_cols_x = 1;
    // int n_cols_y = 1;
    List result_z(3);

    // // check 'n_cols' variable for errors
    // rollshap::check_bounds_int(n_cols_x, 1, 20, "n_cols");
    
    // create a list of matrices
    result_z = roll_lm_z(xx, yy, width,
                         weights, intercept, min_obs,
                         complete_obs, na_restore,
                         online);
    
    arma::vec arma_rsq_z = result_z["r.squared"];
    
    if (intercept) {
      
      // create and return a vector object for Shapley values
      NumericVector rsq(wrap(arma_rsq_z));
      rsq.attr("dim") = IntegerVector::create(n_rows_xy, 1);
      List dimnames_x = xx.attr("dimnames");
      rollshap::xts_attr(rsq, xx, dimnames_x);
      
      return rsq;
      
    } else {
      
      // create and return a vector object for Shapley values
      NumericVector rsq(wrap(arma_rsq_z));
      rollshap::vec_attr(rsq, xx);
      
      return rsq;
      
    }
    
  }
  
}
