// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>
// [[Rcpp::plugins(openmp)]]
#ifdef _OPENMP
  #include <omp.h>
#endif
using namespace Rcpp;

// other librairies:
#include <random> // for mersenne-twister generator
#include <climits> // for UINT_MAX

// Note:
// In your Makevars, use:
// PKG_CXXFLAGS += -O3 -march=native -ffast-math
// PKG_LIBS += -fopenmp
// to improve speed.

// [[Rcpp::export]]
NumericVector devianlm_cpp(const arma::mat & X,
        const arma::uword n_sims,
        const int nthreads) {
  
    const arma::uword n = X.n_rows, k = X.n_cols;
    NumericVector Max_( n_sims );
  
    //--------------------------------------------------
    // Precompute the (already-known, fixed values):
    // hat matrix + scaling factor + various quantities
    //--------------------------------------------------
    // Compute the diagonal of the influence / hat matrix :
    // H = X(X^{T}X)^{-1}X^{T} or in R: X %*% solve(t(X) %*% X) %*% t(X)
    // or for a direct computation one can use: const arma::mat XtX_inv = arma::inv( X.t() * X );
    const arma::mat XtX_inv = arma::solve(
            X.t() * X,
            arma::eye( X.n_cols, X.n_cols),
            arma::solve_opts::likely_sympd );
    const arma::colvec d = arma::sum( X % ( X * XtX_inv ), 1 );
    // Alternatively, one can use arma::sympd() instead of solve():
    // const arma::mat XtX = X.t() * X;
    // const arma::mat XtX_inv = arma::inv_sympd( 0.5 * ( XtX + XtX.t() ) );

    // Also compute 'A':
    const arma::mat A = XtX_inv * X.t(); // (X'X)^{-1} X'

    // Reciprocal of scaling factor:
    const double inv_c = 1.0 / (n - k - 1);

    // 1 - d = 
    const arma::vec one_minus_d = 1.0 - d;

    //--------------------------------------------------
    // Defines RNG and seeds:
    //--------------------------------------------------
    // Define the seed and RNG.
    // one can use std::random_device instead, but it may be slower.
    Rcpp::RNGScope scope;
    // assuming nthreads > 0 (no safeguard here)
    std::vector<unsigned int> seeds( nthreads );
    for( int i = 0; i < nthreads; i++ )
        // span the uniform draw in the [0, UINTMAX[ interval:
        seeds[i] = static_cast<unsigned int>( R::unif_rand() * UINT_MAX );
    
#ifdef _OPENMP
    omp_set_num_threads( nthreads );
#else
    if( nthreads > 1 )
        Rcpp::Rcout << "The OpenMP library is not available on this machine: running in single-threaded mode.\n";
#endif
    //--------------------------------------------------
    // Perform threaded Monte-Carlo simulations:
    //--------------------------------------------------
#ifdef _OPENMP
#pragma omp parallel
#endif
{ // START OF OPEN MP
#ifdef _OPENMP
    // Get thread number and use it as a seed (see below)
    const int thread_id = omp_get_thread_num();
#else
    const int thread_id = 0;
#endif
    
    // seed the rng and initialize the normal distribution,
    // generate entropy from each thread id (2 streams):
    // (see https://en.cppreference.com/w/cpp/numeric/random/seed_seq.html)
    std::seed_seq seq{ seeds[thread_id], (unsigned int)thread_id };
    std::mt19937 rng( seq );
    std::normal_distribution<double> rand_norm( 0.0, 1.0 );    

    // allocate these three vectors once per thread
    //arma::vec y(n), Py(n), resid(n);
    arma::vec y(n), Py(n), resid(n), z(k), w(k), s2i(n), stud_res(n);
    
#ifdef _OPENMP
// we are confident that the runtime cost of each iteration is homogeneous,
// then we use 'static' which distributes the iteration evenly accross the threads.
#pragma omp for schedule(static)
#endif
    for( arma::uword sim_i = 0; sim_i < n_sims; sim_i++ ) {

        // Generate y \sim N(0,1) using Armadillo .imbue() function to populate 'y':
        y.imbue( [&](){ return rand_norm(rng); } );
        // equivalent to:
        // for( arma::uword i = 0; i < n; i++ )
        //      y[i] = rand_norm( rng );

        // equivalent to: y - X*coef; (residuals)
        z = X.t() * y;
        w = XtX_inv * z;
        Py = X * w;
        
        resid = y;
        resid -= Py;

        // Compute residual variance
        const double total_resid_squared = arma::dot(resid, resid); // Dot product for squared sum

        // Compute s2i
        s2i = arma::vec(n);
        s2i.fill(total_resid_squared * inv_c);
        s2i -= (resid % resid) / one_minus_d * inv_c;

        // Compute studentized (external) residuals
        stud_res = resid / arma::sqrt( s2i % one_minus_d );

        // Compute and return the threshold statistic
        Max_[sim_i] = arma::max( arma::abs( stud_res ) );
    }

} // end of OPENMP

    // returns the max of the (external) studentized residuals:
    return Max_;

}
