#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
// 
// Written by:
// -- 
// John L. Weatherwax                2010-02-25
// 
// email: wax@alum.mit.edu
// 
// Please send comments and especially bug reports to the
// above email address.
// 
//-----

using namespace boost::numeric::ublas;
using namespace boost::math;

void ht_regularization( matrix<double> & M, double gamma, size_t num_of_measurements ){
  // H(ypothesis) T(esting) regularization
  //
  // Input:
  //   M                   : A covariance matrix						 
  //   gamma 		   : The amount of regularization. 
  //                         The input gamma in this case is 1-alpha, where alpha is the probability of a
  //                         Type I error, or the probability of rejecting H_0 (the hypothesis that
  //                         \rho_{ij} = 0 or that the population correlation is zero) when H_0 is true thus: 
  //                             gamma ~ 0.0 => alpha = 1.0 => less (no) regularization 
  //                             gamma ~ 1.0 => alpha = 0.0 => maxumum regularization
  //   num_of_measurements : The number of measurements that went into computing this covariance element
  // Output:
  //   M                   : A regularized covariance matrix
  //
  bool verbose = false; 
  double alpha = 1.0 - gamma;
  size_t Sn    = num_of_measurements;
  size_t N     = M.size1(); 

  students_t dist(Sn - 2);
  double t = quantile(complement(dist, alpha/2));

  for( size_t ii=0; ii<N; ++ii ){
    for( size_t jj=0; jj<N; ++jj ){
      if( ii==jj ) continue;
      double r = M(ii,jj)/std::sqrt( M(ii,ii)*M(jj,jj) );
      double T = r * std::sqrt( (Sn - 2) / (1-r*r) );
      if( verbose ) std::cout << "(ii,jj)=( " << std::setw(6) << ii << ", " << std::setw(6) << jj << ") T= " << std::setw(10) << T << " t=" << std::setw(10) << t << " n=" << std::setw(6) << Sn;
      if( T < t ){
	if( verbose ) std::cout << " assigning to zero ... " << std::endl; 
	M(ii,jj) = 0.;                 // modify non-diagonal elements 
      }else{
	if( verbose ) std::cout << std::endl;
      }
    }
  }
}