// #include <Rcpp.h>
#include <RcppArmadillo.h>
#include "beachmat/numeric_matrix.h"
#include "beachmat/integer_matrix.h"

#include <deviance.h>
#include <fisher_scoring_steps.h>

using namespace Rcpp;



template<class NumericType>
void clamp_inplace(/*INOUT parameter*/ arma::Mat<NumericType>& v, double min, double max){
  for(int i = 0; i < v.n_elem; i++){
    if(v.at(i) < min){
      v.at(i) = min;
    }else if(v.at(i) > max){
      v.at(i) = max;
    }
  }
}


// Check how many unique rows are in a matrix and if this number is less than or equal to n
// This is important to determine if the model can be solved by group averages
// (ie. the numer of unique rows == number of columns)
// [[Rcpp::export]]
bool lte_n_equal_rows(const NumericMatrix& matrix, int n, double tolerance = 1e-10) {
  NumericMatrix reference_matrix(n, matrix.ncol());
  size_t n_matches = 0;
  for(size_t row_idx = 0; row_idx < matrix.nrow(); row_idx++){
    bool matched = false;
    NumericMatrix::ConstRow vec = matrix(row_idx, _);
    for(size_t ref_idx = 0; ref_idx < n_matches; ref_idx++){
      NumericMatrix::Row ref_vec  = reference_matrix(ref_idx, _);
      if(sum(abs(vec - ref_vec)) < tolerance){
        matched = true;
        break;
      }
    }
    if(! matched){
      ++n_matches;
      if(n_matches > n){
        return false;
      }
      reference_matrix(n_matches - 1, _) = vec;
    }
  }
  return true;
}

// [[Rcpp::export]]
IntegerVector get_row_groups(const NumericMatrix& matrix, int n_groups, double tolerance = 1e-10) {
  NumericMatrix reference_matrix(n_groups, matrix.ncol());
  IntegerVector groups(matrix.nrow());
  size_t n_matches = 0;
  for(size_t row_idx = 0; row_idx < matrix.nrow(); row_idx++){
    bool matched = false;
    NumericMatrix::ConstRow vec = matrix(row_idx, _);
    for(size_t ref_idx = 0; ref_idx < n_matches; ref_idx++){
      NumericMatrix::Row ref_vec  = reference_matrix(ref_idx, _);
      if(sum(abs(vec - ref_vec)) < tolerance){
        groups(row_idx) = ref_idx;
        matched = true;
        break;
      }
    }
    if(! matched){
      groups(row_idx) = n_matches;
      reference_matrix(n_matches, _) = vec;
      ++n_matches;
    }
  }
  return groups + 1;
}


arma::vec calculate_mu(const arma::mat& model_matrix, const arma::vec& beta_hat, const arma::vec& exp_off){
  arma::vec mu_hat = exp(model_matrix * beta_hat) % exp_off;
  clamp_inplace(mu_hat, 1e-50, 1e50);
  return mu_hat;
}

/**
 * This method takes in a proposal for a step and checks if it actually
 * decreases the deviance of the model. If does not add it tries again
 * with half the step size, then a quarter and so on.
 *
 * If even after 100 steps the deviance (0.5^100 = 7.9e-31) has not decreased
 * it returns NaN.
 *
 * Note that the first two parameters are changed: beta_hat and mu_hat
 *
 * The function returns the new deviance.
 *
 */
template<class NumericType>
double decrease_deviance(/*In-Out Parameter*/ arma::vec& beta_hat,
                         /*In-Out Parameter*/ arma::vec& mu_hat,
                         const arma::vec& step,
                         const arma::mat& model_matrix,
                         const arma::mat& exp_off,
                         const arma::Col<NumericType>& counts,
                         const double theta, const double dev_old, const double tolerance, const double max_rel_mu_change){
  double speeding_factor = 1.0;
  int line_iter = 0;
  double dev = 0;
  beta_hat = beta_hat + step;
  const arma::vec mu_old = mu_hat;
  while(true){
    mu_hat = calculate_mu(model_matrix, beta_hat, exp_off);
    dev = compute_gp_deviance_sum(counts, mu_hat, theta);
    double conv_test = fabs(dev - dev_old)/(fabs(dev) + 0.1);
    double mu_rel_change = max(mu_hat / mu_old);
    if((dev < dev_old && mu_rel_change < max_rel_mu_change) || conv_test < tolerance){
      break; // while loop
    }else if(line_iter >= 100){
      // speeding factor is very small, something is going wrong here
      dev = std::numeric_limits<double>::quiet_NaN();
      break; // while loop
    }else{
      // Halfing the speed
      speeding_factor = speeding_factor / 2.0;
      beta_hat = beta_hat - step * speeding_factor;
    }
    line_iter++;
  }
  return dev;
}


template<class NumericType>
double decrease_deviance_plus_ridge(/*In-Out Parameter*/ arma::vec& beta_hat,
                                    /*In-Out Parameter*/ arma::vec& mu_hat,
                                    const arma::vec& step,
                                    const arma::mat& model_matrix,
                                    const arma::mat& ridge_penalty_sq,
                                    const arma::vec& ridge_target,
                                    const arma::mat& exp_off,
                                    const arma::Col<NumericType>& counts,
                                    const double theta, const double dev_old,
                                    const double tolerance, const double max_rel_mu_change){
  double speeding_factor = 1.0;
  int line_iter = 0;
  double dev = 0;
  int n_samples = model_matrix.n_rows;
  beta_hat = beta_hat + step;
  const arma::vec mu_old = mu_hat;
  while(true){
    mu_hat = calculate_mu(model_matrix, beta_hat, exp_off);
    double pen_sum = n_samples * arma::as_scalar((beta_hat - ridge_target).t() * ridge_penalty_sq * (beta_hat - ridge_target));
    dev = compute_gp_deviance_sum(counts, mu_hat, theta) + pen_sum;
    double conv_test = fabs(dev - dev_old)/(fabs(dev) + 0.1);
    double mu_rel_change = max(mu_hat / mu_old);
    if((dev < dev_old && mu_rel_change < max_rel_mu_change) || conv_test < tolerance){
      break; // while loop
    }else if(line_iter >= 100){
      // speeding factor is very small, something is going wrong here
      dev = std::numeric_limits<double>::quiet_NaN();
      break; // while loop
    }else{
      // Halfing the speed
      speeding_factor = speeding_factor / 2.0;
      beta_hat = beta_hat - step * speeding_factor;
    }
    line_iter++;
  }
  return dev;
}




//--------------------------------------------------------------------------------------------------//
// The following code was originally copied from https://github.com/mikelove/DESeq2/blob/master/src/DESeq2.cpp
// I adapted it to the needs of this project by:
//  * remove weights
//  * Calculate actual deviance (2 * (log(f_NB(y | mu, theta)) - log(f_NB(y | y, theta))))
//    instead of just 2 * log(f_NB(y | mu, theta)),
//  * Support DelayedArrays
//  * Remove unncessary outputs: beta_mat_var, hat_diagonals, deviance
//  * Remove beta divergence check if abs(beta) very large
//  * Add line search that ensures that deviance is decreasing at every step
//  * Add optional ridge penalty


// fit the Negative Binomial GLM with Fisher scoring
// note: the betas are on the natural log scale
//
template<class NumericType, class BMNumericType>
List fitBeta_fisher_scoring_impl(RObject Y, const arma::mat& model_matrix, RObject exp_offset_matrix,
                                 NumericVector thetas, SEXP beta_matSEXP, Nullable<NumericMatrix> ridge_penalty_nl,
                                 double tolerance, double max_rel_mu_change, int max_iter, bool use_diagonal_approx) {
  auto Y_bm = beachmat::create_matrix<BMNumericType>(Y);
  auto exp_offsets_bm = beachmat::create_numeric_matrix(exp_offset_matrix);
  int n_samples = Y_bm->get_ncol();
  int n_genes = Y_bm->get_nrow();

  // the ridge penalty
  bool apply_ridge_penalty = ridge_penalty_nl.isNotNull();
  arma::mat ridge_penalty;
  arma::mat ridge_penalty_sq;
  arma::vec ridge_target;
  if(apply_ridge_penalty){
    NumericMatrix tmp = ridge_penalty_nl.get();
    ridge_penalty = arma::mat(tmp.cbegin(), tmp.nrow(), tmp.ncol());
    if(model_matrix.n_cols != ridge_penalty.n_cols){
      stop("Number of columns in model_matrix does not match the columns of the ridge_penalty");
    }
    ridge_penalty_sq = ridge_penalty.t() * ridge_penalty;

    if(tmp.hasAttribute("target")){
      ridge_target = (NumericVector) tmp.attr("target");
    }else{
      ridge_target = arma::zeros(tmp.ncol());
    }
  }
  // The result
  arma::mat beta_mat = as<arma::mat>(beta_matSEXP);

  // deviance, convergence and tolerance
  NumericVector iterations(n_genes);
  NumericVector deviance(n_genes);
  for (int gene_idx = 0; gene_idx < n_genes; gene_idx++) {
    if (gene_idx % 100 == 0) checkUserInterrupt();
    // Fill count and offset vector from beachmat matrix
    arma::Col<NumericType> counts(n_samples);
    Y_bm->get_row(gene_idx, counts.begin());
    arma::Col<double> exp_off(n_samples);
    exp_offsets_bm->get_row(gene_idx, exp_off.begin());
    // Init beta and mu
    arma::vec beta_hat = beta_mat.row(gene_idx).t();
    arma::vec mu_hat = calculate_mu(model_matrix, beta_hat, exp_off);
    if(beta_hat.has_nan() || Rcpp::traits::is_na<REALSXP>(thetas(gene_idx))){
      beta_hat.fill(NA_REAL);
      iterations(gene_idx) = 0;
      deviance(gene_idx) = NA_REAL;
      continue;
    }
    // Init deviance
    double dev_old = 0;
    if(apply_ridge_penalty){
      // For diagonal ridge_penalty: pen = Sum (lambda_i b_i)^2
      double pen_sum = n_samples * arma::as_scalar((beta_hat - ridge_target).t() * ridge_penalty_sq * (beta_hat - ridge_target));
      dev_old = compute_gp_deviance_sum(counts, mu_hat, thetas(gene_idx)) + pen_sum;
    }else{
      dev_old = compute_gp_deviance_sum(counts, mu_hat, thetas(gene_idx));
    }
    for (int t = 0; t < max_iter; t++) {
      iterations(gene_idx)++;
      // Find good direction to optimize beta
      arma::vec step;
      if(use_diagonal_approx){
        step = fisher_scoring_diagonal_step(model_matrix, counts, mu_hat, thetas(gene_idx) * mu_hat);
      }else{
        if(apply_ridge_penalty){
          step = fisher_scoring_qr_ridge_step(model_matrix, counts, mu_hat, thetas(gene_idx) * mu_hat, ridge_penalty, ridge_target, beta_hat);
        }else{
          step = fisher_scoring_qr_step(model_matrix, counts, mu_hat, thetas(gene_idx) * mu_hat);
        }
      }
      // Find step size that actually decreases the deviance
      double dev = 0;
      if(apply_ridge_penalty){
        dev = decrease_deviance_plus_ridge(beta_hat, mu_hat, step, model_matrix, ridge_penalty_sq, ridge_target,
                                exp_off, counts, thetas(gene_idx), dev_old, tolerance, max_rel_mu_change);
      }else{
        dev = decrease_deviance(beta_hat, mu_hat, step, model_matrix,
                                exp_off, counts, thetas(gene_idx), dev_old, tolerance, max_rel_mu_change);
      }
      double conv_test = fabs(dev - dev_old)/(fabs(dev) + 0.1);
      dev_old = dev;
      if (std::isnan(conv_test)) {
        // This should not happen
        beta_hat.fill(NA_REAL);
        iterations(gene_idx) = max_iter;
        break;
      }
      if (conv_test < tolerance) {
        break;
      }
    }
    beta_mat.row(gene_idx) = beta_hat.t();
    deviance(gene_idx) = dev_old;
  }

  return List::create(
    Named("beta_mat", beta_mat),
    Named("iter", iterations),
    Named("deviance", deviance));
}


// [[Rcpp::export]]
List fitBeta_fisher_scoring(RObject Y, const arma::mat& model_matrix, RObject exp_offset_matrix,
                                  NumericVector thetas, SEXP beta_matSEXP, Nullable<NumericMatrix> ridge_penalty_nl,
                                  double tolerance, double max_rel_mu_change, int max_iter) {
  auto mattype=beachmat::find_sexp_type(Y);
  if (mattype==INTSXP) {
    return fitBeta_fisher_scoring_impl<int, beachmat::integer_matrix>(Y, model_matrix, exp_offset_matrix,
                                                                      thetas,  beta_matSEXP,
                                                                      /*ridge_penalty=*/ ridge_penalty_nl,
                                                                      tolerance, max_rel_mu_change, max_iter,
                                                                      /*use_diagonal_approx=*/ false);
  } else if (mattype==REALSXP) {
    return fitBeta_fisher_scoring_impl<double, beachmat::numeric_matrix>(Y, model_matrix, exp_offset_matrix,
                                                                         thetas,  beta_matSEXP,
                                                                         /*ridge_penalty=*/ ridge_penalty_nl,
                                                                         tolerance, max_rel_mu_change, max_iter,
                                                                         /*use_diagonal_approx=*/ false);
  } else {
    throw std::runtime_error("unacceptable matrix type");
  }
}



// [[Rcpp::export]]
List fitBeta_diagonal_fisher_scoring(RObject Y, const arma::mat& model_matrix, RObject exp_offset_matrix,
                                     NumericVector thetas, SEXP beta_matSEXP,
                                     double tolerance, double max_rel_mu_change, int max_iter) {
  auto mattype=beachmat::find_sexp_type(Y);
  if (mattype==INTSXP) {
    return fitBeta_fisher_scoring_impl<int, beachmat::integer_matrix>(Y, model_matrix, exp_offset_matrix,
                                                                      thetas,  beta_matSEXP,
                                                                      /*ridge_penalty=*/ R_NilValue,
                                                                      tolerance, max_rel_mu_change, max_iter,
                                                                      /*use_diagonal_approx=*/ true);
  } else if (mattype==REALSXP) {
    return fitBeta_fisher_scoring_impl<double, beachmat::numeric_matrix>(Y, model_matrix, exp_offset_matrix,
                                                                         thetas,  beta_matSEXP,
                                                                         /*ridge_penalty=*/ R_NilValue,
                                                                         tolerance, max_rel_mu_change, max_iter,
                                                                         /*use_diagonal_approx=*/ true);
  } else {
    throw std::runtime_error("unacceptable matrix type");
  }
}



// If there is only one group, there is no need to do the full Fisher-scoring
// Instead a simple Newton-Raphson algorithm will do
template<class NumericType>
List fitBeta_one_group_internal(SEXP Y_SEXP, SEXP offsets_SEXP,
                       NumericVector thetas, NumericVector beta_start_values,
                       double tolerance, int maxIter) {
  auto Y_bm = beachmat::create_matrix<NumericType>(Y_SEXP);

  auto offsets_bm = beachmat::create_numeric_matrix(offsets_SEXP);
  int n_samples = Y_bm->get_ncol();
  int n_genes = Y_bm->get_nrow();
  NumericVector result(n_genes);
  IntegerVector iterations(n_genes);
  NumericVector deviance(n_genes);

  Environment glmGamPoiEnv = Environment::namespace_env("glmGamPoi");
  Function estimate_betas_group_wise_optimize_helper = glmGamPoiEnv["estimate_betas_group_wise_optimize_helper"];

  for(int gene_idx = 0; gene_idx < n_genes; gene_idx++){
    if (gene_idx % 100 == 0) checkUserInterrupt();

    double beta = beta_start_values(gene_idx);
    const double& theta = thetas(gene_idx);
    if(Rcpp::traits::is_na<REALSXP>(beta) || Rcpp::traits::is_na<REALSXP>(theta)){
      // Missing values, just continue with next gene
      result(gene_idx) = NA_REAL;
      iterations(gene_idx) = 0;
      deviance(gene_idx) = NA_REAL;
      continue;
    }

    typename NumericType::vector counts(n_samples);
    Y_bm->get_row(gene_idx, counts.begin());
    NumericVector off(n_samples);
    offsets_bm->get_row(gene_idx, off.begin());
    // Newton-Raphson
    int iter = 0;
    for(; iter < maxIter; iter++){
      double dl = 0.0;
      double ddl = 0.0;
      bool all_zero = true;
      for(int sample_iter = 0; sample_iter < n_samples; sample_iter++){
        const auto count = counts[sample_iter];
        all_zero = all_zero && count == 0;
        const double mu = std::exp(beta + off[sample_iter]);
        const double denom = 1.0 + mu * theta;
        dl += (count - mu) / denom;
        ddl += mu * (1.0 + count * theta) / denom / denom;
        // ddl += mu / denom;           // This is what edgeR is using
      }
      if(all_zero){
        beta = R_NegInf;
        break;
      }
      const double step = dl / ddl;
      beta += step;
      if(std::abs(step) < tolerance){
        break;
      }else if(Rcpp::traits::is_nan<REALSXP>(beta)){
        break;
      }
    }
    if(iter == maxIter || Rcpp::traits::is_nan<REALSXP>(beta)){
      // Not converged -> try again with optimize()
      beta =  Rcpp::as<double>(estimate_betas_group_wise_optimize_helper(counts, off, theta));
    }
    result(gene_idx) = beta;
    iterations(gene_idx) = iter;
    double dev = 0.0;
    for(int sample_iter = 0; sample_iter < n_samples; sample_iter++){
      dev += compute_gp_deviance(counts[sample_iter], exp(beta + off[sample_iter]), theta);
    }
    deviance(gene_idx) = dev;
  }
  return List::create(
    Named("beta", result),
    Named("iter", iterations),
    Named("deviance", deviance)
  );
}

// [[Rcpp::export(rng = false)]]
List fitBeta_one_group(RObject Y, RObject offset_matrix,
                        NumericVector thetas, NumericVector beta_start_values,
                        double tolerance, int maxIter) {
  auto mattype=beachmat::find_sexp_type(Y);
  if (mattype==INTSXP) {
    return fitBeta_one_group_internal<beachmat::integer_matrix>(Y, offset_matrix, thetas, beta_start_values, tolerance, maxIter);
  } else if (mattype==REALSXP) {
    return fitBeta_one_group_internal<beachmat::numeric_matrix>(Y, offset_matrix, thetas, beta_start_values, tolerance, maxIter);
  } else {
    throw std::runtime_error("unacceptable matrix type");
  }
}









