#include <RcppArmadillo.h>
#include <cmath>
// [[Rcpp::plugins(cpp17)]]
// [[Rcpp::depends(RcppArmadillo)]]
using namespace Rcpp;
using namespace arma;

// wrapper around R's RNG such that we get a uniform distribution over
// [0,n) as required by the STL algorithm
//inline int randWrapper(const int n) { return floor(unif_rand()*n); }

unsigned int perm_one_gene(arma::vec const& breaks, 
                           arma::mat& cdf_A, 
                           arma::mat& cdf_B,
                           unsigned int const& n_samples,
                           arma::uvec const& sample_has_cells,
                           arma::uvec const& sample_ids_one_cl,
                           arma::vec const& counts_one_gene,
                           arma::vec const& group_ids_of_samples,
                           unsigned int const& N_breaks,
                           unsigned int const& n_samples_A,
                           unsigned int const& n_samples_B,
                           arma::vec& T_perm,
                           unsigned int const& P,
                           unsigned int const& n_both,
                           arma::umat const& PERM,
                           double T_obs)
{
  unsigned int res = 0;
  unsigned int tmp2, cell, b, sample, p;
  
  /* Permutation testing */
  T_perm.fill(0.0);
  arma::vec counts_one_gene_permuted(n_both);
  
  for (p=0 ; p<P ; p++) {
    
    for (b=0 ; b<n_both ; b++) {
      counts_one_gene_permuted(b) = counts_one_gene( PERM(b,p) );
    }
    
    cdf_A.fill(0.0);
    cdf_B.fill(0.0);
    for (sample=0 ; sample<n_samples ; sample++ ) {
      if ( (sample_has_cells(sample))==1 ) { // check that samples have at least 1 cell
        arma::uvec tmp1 = find( sample_ids_one_cl == sample ); // find the cells associated to the sample `sample`
        tmp2 = tmp1.n_elem; // counts number of cells
        arma::vec counts_one_gene_one_sample(tmp2); // create a vector to store these cells
        
        for (b=0 ; b<tmp2 ; b++) { // store cells of sample `sample`
          counts_one_gene_one_sample(b) = counts_one_gene_permuted( tmp1(b) );
        }
        
        // sort counts in ascending order
        std::sort(counts_one_gene_one_sample.begin(), counts_one_gene_one_sample.end());
        
        cell = 0; // cell index
        
        if ( (group_ids_of_samples(sample)) == 1.0) { // if sample belongs to group A, add to cdf_A
          
          for (b=0 ; b<N_breaks ; b++) { // loop over the number of breaks of the CDF
            // loop until cell index is smaller than tmp2 (length of counts_one_gene_one_sample)
            while(cell < tmp2){
              if (breaks(b)>counts_one_gene_one_sample(cell) ) { // if break is > cell, keep counting and go to next cell
                cell += 1;
              }
              else{ // if break is < cell, stop loop and save cdf for break b
                cdf_A(b) += 1.0 * cell/(tmp2*n_samples_A);
                break;
              }
            }
            
            if(cell == tmp2){
              cdf_A(b) += 1.0/n_samples_A;
            }
            
          } //end of b for loop
        } // end of GROUP 1
        else { // if sample belongs to group B, add to cdf_B
          
          for (b=0 ; b<N_breaks ; b++) { // loop over the number of breaks of the CDF
            // loop until cell index is smaller than tmp2 (length of counts_one_gene_one_sample)
            while(cell < tmp2){
              if (breaks(b)>counts_one_gene_one_sample(cell) ) { // if break is > cell, keep counting and go to next cell
                cell += 1;
              }
              else{ // if break is < cell, stop loop and save cdf for break b
                cdf_B(b) += 1.0 * cell/(tmp2*n_samples_B);
                break;
              }
            }
            
            if(cell == tmp2){
              cdf_B(b) += 1.0/n_samples_B;
            }
            
          } //end of b for loop
        } // end of GROUP 2
        
      }
    }
    
    /* Permuted test statistic */
    for (b=0 ; b<N_breaks ; b++) {
      T_perm(p) += fabs( cdf_A(b) - cdf_B(b) );
    }
    
  }   // End of loop for permutations
  
  for (p=0 ; p<P ; ++p) {
    if (T_perm(p)>=T_obs) {
      res += 1;
    }
  }
  
  return res;
}


// [[Rcpp::export]]
List perm_test(unsigned int const& P,                             // number of permutations for all gene-cluster combinations
               unsigned int const& P_2,                           // number of permutations when p < 0.1
               unsigned int const& P_3,                           // number of permutations when p < 0.01
               unsigned int const& P_4,                           // number of permutations when p < 0.001
               unsigned int const& N_breaks,                      // number of breaks at which to evaluate the cdf
               arma::vec const& cluster_ids,                            // ids of clusters (cell-population) for every cell
               unsigned int const& n_clusters,                    // total number of clusters
               arma::vec const& sample_ids,                             // ids of samples for every cell
               unsigned int const& n_samples,                     // total number of samples
               arma::vec const& group_ids_of_samples,                   // ids of groups (1 or 2) for every sample
               double const& min_non_zero_cells,                  // min number of cells with > 0 expression in each cluster to consider the gene for testing
               arma::sp_mat const& counts)                                // count matrix (rows = genes, cols = cells)
{
  arma::vec T_perm(P), T_perm_2(P_2 - P), T_perm_3(P_3 - P_2), T_perm_4(P_4 - P_3);
  
  /* Global variables */
  unsigned int tmp2, cell, b, sample, p, n_cells;
  
  double T_obs, p_val;
  
  unsigned int n_genes = counts.n_cols, tmp_res; //, n_cells_overall = counts.n_rows;
  arma::uvec samples_in_group_A = arma::find( group_ids_of_samples == 1);
  arma::uvec samples_in_group_B = arma::find( group_ids_of_samples == 2);
  double s, mn, len_out;
  unsigned int n_samples_A = samples_in_group_A.n_elem;                        // number of samples in group A
  unsigned int n_samples_B = samples_in_group_B.n_elem;                        // number of samples in group B
  arma::mat res(n_genes,n_clusters); res.fill(-1);                             // array of p-values
  arma::vec seq0(N_breaks), breaks(N_breaks);
  for (b=0 ; b<N_breaks ; b++) {
    seq0(b) = b + 1.0;
  }
  
  arma::vec cdf_A(N_breaks);
  arma::vec cdf_B(N_breaks); 
  
  /* Start of cluster loop - declare the variables whose size won't change */
  arma::uvec sample_has_cells(n_samples);
  for (unsigned int cl_id = 0; cl_id < n_clusters; ++cl_id) { // loop over clusters
    // select data belonging to both groups:
    arma::uvec ids_both = arma::find( cluster_ids == cl_id );
    unsigned int n_both = ids_both.n_elem;
    arma::vec counts_one_gene(n_both);
    
    // vector with data for one gene (all cells of all clusters):
    //arma::vec counts_gene(n_cells_overall);
    
    n_cells = ids_both.n_elem;
    arma::uvec sample_ids_one_cl(n_cells); 
    for (b=0 ; b<n_cells ; b++) {
      sample_ids_one_cl(b) = sample_ids( ids_both(b) );
    }
    
    /* permutation martix */ 
    arma::umat PERM(n_cells,P);
    Rcpp::IntegerMatrix PERMUTATIONS(n_cells,P);
    Rcpp::IntegerVector permutation = Rcpp::seq_len(n_cells)-1; // define the vector for each permutation, -1 because Rcpp indexes start from 0 !
    for (p=0 ; p<P ; p++) {
      PERM.col(p) = randperm(n_cells);
      //std::random_shuffle(permutation.begin(), permutation.end(), randWrapper);
      //PERMUTATIONS(_,p) = permutation;
    }
    PERM = Rcpp::as<arma::umat>(PERMUTATIONS);
    
    /* SECOND permutation martix */ 
    arma::umat PERM_2(n_cells,P_2-P);
    Rcpp::IntegerMatrix PERMUTATIONS_2(n_cells,P_2-P);
    for (p=0 ; p<(P_2-P) ; p++) {
      PERM_2.col(p) = randperm(n_cells);
      //std::random_shuffle(permutation.begin(), permutation.end(), randWrapper);
      //PERMUTATIONS_2(_,p) = permutation;
    }
    // PERM_2 = Rcpp::as<arma::umat>(PERMUTATIONS_2);
    
    /* THIRD permutation martix */ 
    arma::umat PERM_3(n_cells,P_3-P_2);
    Rcpp::IntegerMatrix PERMUTATIONS_3(n_cells,P_3-P_2);
    for (p=0 ; p<(P_3-P_2) ; p++) {
      PERM_3.col(p) = randperm(n_cells);
      //std::random_shuffle(permutation.begin(), permutation.end(), randWrapper);
      //PERMUTATIONS_3(_,p) = permutation;
    }
    //PERM_3 = Rcpp::as<arma::umat>(PERMUTATIONS_3);
    
    /* THIRD permutation martix */ 
    arma::umat PERM_4(n_cells,P_4-P_3);
    Rcpp::IntegerMatrix PERMUTATIONS_4(n_cells,P_4-P_3);
    for (p=0 ; p<(P_4-P_3) ; p++) {
      PERM_4.col(p) = randperm(n_cells);
      // std::random_shuffle(permutation.begin(), permutation.end(), randWrapper);
      // PERMUTATIONS_4(_,p) = permutation;
    }
    //PERM_4 = Rcpp::as<arma::umat>(PERMUTATIONS_4);
    
    sample_has_cells.fill(0);
    for (sample=0 ; sample<n_samples ; sample++) {
      sample_has_cells(sample) = any(sample_ids_one_cl==sample);
    }
    
    /* Flag to interrupt the code from R: */
    Rcpp::checkUserInterrupt();
    
    /* gene loop */ 
    for (unsigned int gene = 0; gene < n_genes; ++gene) {
      //Rcpp::checkUserInterrupt();
      
      /* Determine if there are sufficient counts */
      //counts_gene = counts.col(gene);
      arma::sp_mat counts_gene(counts.col(gene));
      
      s = 0.0;
      for (b=0 ; b<n_both ; b++) {
        counts_one_gene(b) = counts_gene( ids_both(b) );
        if ( counts_one_gene(b) > 0 ) {
          s += 1.0;
        }
      } 
      
      /* */
      if (s >= min_non_zero_cells) {
        // if keep_gene is false (0), res(gene,cl_id) will be -1 (filled above).
        // if keep_gene is true (1),res(gene,cl_id) will be set at 1 (the observed value) and values added below (and then devided by P+1).
        
        /* CDF grid */ 
        mn = counts_one_gene.min();
        len_out = ( counts_one_gene.max() - mn ) / (N_breaks + 1.0);
        for (b=0 ; b<N_breaks ; b++) {
          breaks(b) = mn + len_out*seq0(b) ;
        }
        
        /* compute T_obs */ 
        cdf_A.fill(0.0);
        cdf_B.fill(0.0);
        for ( sample=0 ; sample<n_samples ; sample++ ) {
          if ( (sample_has_cells(sample))==1 ) { // check that samples have at least 1 cell
            arma::uvec tmp1 = find( sample_ids_one_cl == sample ); // find the cells associated to the sample `sample`
            tmp2 = tmp1.n_elem; // counts number of cells
            arma::vec counts_one_gene_one_sample(tmp2); // create a vector to store these cells
            for ( b=0 ; b<tmp2 ; b++) { // store cells of sample `sample`
              counts_one_gene_one_sample(b) = counts_one_gene( tmp1(b) );
            }
            
            // sort counts in ascending order
            std::sort(counts_one_gene_one_sample.begin(), counts_one_gene_one_sample.end());
            
            cell = 0; // cell index
            
            if ( (group_ids_of_samples(sample)) == 1.0) { // if sample belongs to group A, add to cdf_A
              
              for ( b=0 ; b<N_breaks ; b++) { // loop over the number of breaks of the CDF
                // loop until cell index is smaller than tmp2 (length of counts_one_gene_one_sample)
                while(cell < tmp2){
                  if (breaks(b)>counts_one_gene_one_sample(cell) ) { // if break is > cell, keep counting and go to next cell
                    cell += 1;
                  }
                  else{ // if break is < cell, stop loop and save cdf for break b
                    cdf_A(b) += 1.0 * cell/(tmp2*n_samples_A);
                    break;
                  }
                }
                
                if(cell == tmp2){
                  cdf_A(b) += 1.0/n_samples_A;
                }
                
              } //end of b for loop
            } // end of GROUP 1
            else { // if sample belongs to group B, add to cdf_B
              
              for ( b=0 ; b<N_breaks ; b++) { // loop over the number of breaks of the CDF
                // loop until cell index is smaller than tmp2 (length of counts_one_gene_one_sample)
                while(cell < tmp2){
                  if (breaks(b)>counts_one_gene_one_sample(cell) ) { // if break is > cell, keep counting and go to next cell
                    cell += 1;
                  }
                  else{ // if break is < cell, stop loop and save cdf for break b
                    cdf_B(b) += 1.0 * cell/(tmp2*n_samples_B);
                    break;
                  }
                }
                
                if(cell == tmp2){
                  cdf_B(b) += 1.0/n_samples_B;
                }
                
              } //end of b for loop
            } // end of GROUP 2
            
          }
        }
        
        /* Obseved test statistic */
        T_obs = 0.0;
        for ( b=0 ; b<N_breaks ; b++) {
          T_obs += fabs( cdf_A(b) - cdf_B(b) )  ;
        }
        
        // run 100 permutations on all gene-cluster combinations
        tmp_res = 1 + perm_one_gene(breaks, 
                                    cdf_A, 
                                    cdf_B,
                                    n_samples,
                                    sample_has_cells,
                                    sample_ids_one_cl,
                                    counts_one_gene,
                                    group_ids_of_samples,
                                    N_breaks,
                                    n_samples_A,
                                    n_samples_B,
                                    T_perm,
                                    P,
                                    n_both,
                                    PERM,
                                    T_obs);
        
        p_val = (1.0 * tmp_res)/( P + 1.0 );
        
        if( p_val <= 0.1 ){ // if p_val < 0.1, use 10 * perm
          if(P_2 > P){
            tmp_res += perm_one_gene(breaks, 
                                     cdf_A, 
                                     cdf_B,
                                     n_samples,
                                     sample_has_cells,
                                     sample_ids_one_cl,
                                     counts_one_gene,
                                     group_ids_of_samples,
                                     N_breaks,
                                     n_samples_A,
                                     n_samples_B,
                                     T_perm_2, // 2nd T_perm_2
                                     P_2 - P,  // 2nd P = (P_2 - P)
                                     n_both,
                                     PERM_2,    // 2nd PERM
                                     T_obs);
            
            p_val = (1.0 * tmp_res)/( P_2 + 1.0 );
          }
          
          if( p_val <= 0.01 ){ // if p_val < 0.01, use 100 * perm
            if(P_3 > P_2){
              tmp_res += perm_one_gene(breaks, 
                                       cdf_A, 
                                       cdf_B,
                                       n_samples,
                                       sample_has_cells,
                                       sample_ids_one_cl,
                                       counts_one_gene,
                                       group_ids_of_samples,
                                       N_breaks,
                                       n_samples_A,
                                       n_samples_B,
                                       T_perm_3,  // 3rd T_perm_3
                                       P_3 - P_2, // 3rd P = P_3-(P_2)
                                       n_both,
                                       PERM_3,    // 2nd PERM
                                       T_obs);
              
              p_val = (1.0 * tmp_res)/( P_3 + 1.0 );
            }
            
            if( p_val <= 0.001 ){ // if p_val < 0.001, use 1,000 * perm
              if(P_4 > P_3){
                tmp_res += perm_one_gene(breaks, 
                                         cdf_A, 
                                         cdf_B,
                                         n_samples,
                                         sample_has_cells,
                                         sample_ids_one_cl,
                                         counts_one_gene,
                                         group_ids_of_samples,
                                         N_breaks,
                                         n_samples_A,
                                         n_samples_B,
                                         T_perm_4,  // 3rd T_perm_3
                                         P_4 - P_3, // 3rd P = P_3-(P_2)
                                         n_both,
                                         PERM_4,    // 2nd PERM
                                         T_obs);
                
                p_val = (1.0 * tmp_res)/( P_4 + 1.0 );
              }
            }
          }
        }
        
        res(gene,cl_id) = p_val;
      }
    }       // End of loop for genes
  }         // End of loop for clusters
  
  List output = List::create( res );
  return output;
}
