#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.
// 
//-----

void covMarket( matrix<double> & x, double shrink, matrix<double> & sigma ){
  //
  // Input: 
  //    x (t*n): t iid observations on n random variables
  //    shrink : shrinkage amount
  // Output: 
  //    sigma (n*n): invertible covariance matrix estimator
  // 
  // This estimator is a weighted average of the sample
  // covariance matrix and a "prior" or "shrinkage target".
  // Here, the prior is given by a one-factor model.
  // The factor is equal to the cross-sectional average
  // of all the random variables.
  //
  // The notation follows Ledoit and Wolf (2003)
  //
  assert( (0.0 <= shrink) && (shrink <= 1.0) );

  bool demeanReturns = true; 

  size_t t=x.size1();
  size_t n=x.size2();

  // Do we demean the stock returns: 
  if( demeanReturns ){
      // compute the mean of the returns: 
      matrix<double> meanx(1,n); meanx.clear();
      for( size_t jj=0; jj<n; ++jj ){
	for( size_t ii=0; ii<t; ++ii ){
	  meanx(0,jj) += x(ii,jj)/t; 
	}
      }

      // de-mean the returns:
      for( size_t jj=0; jj<n; ++jj ){
	for( size_t ii=0; ii<t; ++ii ){
	  x(ii,jj) -= meanx(0,jj);
	}
      }
  }

  // compute a equally weighted average market return vector "xmkt":
  matrix<double> xmkt(t,1); xmkt.clear();
  for( size_t ii=0; ii<t; ++ii ){
    for( size_t jj=0; jj<n; ++jj ){
      xmkt(ii,0) += x(ii,jj)/n;
    }
  }
  // compute the variance of the market "varmkt":
  double varmkt=0.;
  for( size_t ii=0; ii<t; ++ii ) varmkt += xmkt(ii,0)*xmkt(ii,0)/t;
  
  // compute the sample covariance matrix:
  matrix<double> sample(n,n);
  sample = prod( trans(x), x ) / t;
  
  // compute the covariance of xmkt with each stock "covmkt":
  matrix<double> covmkt(n,1); covmkt.clear();
  for( size_t ii=0; ii<t; ++ii ){
    for( size_t jj=0; jj<n; ++jj ){
      covmkt(jj,0) += x(ii,jj)*xmkt(ii,0)/t;
    }
  }

  // This line is equivalent to s_00^2 b b' in section 2.3 of the above paper: 
  matrix<double> prior(n,n); prior.clear();
  prior  = prod( covmkt, trans(covmkt) ) / varmkt;

  // override the diagonal elements: 
  for( size_t ii=0; ii<n; ++ii ) prior(ii,ii) = sample(ii,ii);

  // compute the estimator:
  sigma.resize(n,n); 
  sigma = shrink * prior + ( 1. - shrink ) * sample;
}