#include "functions.hpp"
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>

namespace gd {

using namespace boost::numeric::ublas;

void py_export_functions() {
	using namespace boost::python;
	def("matmul_2d_3d", matmul_2d_3d);
	def("matmul_3d_3d", matmul_3d_3d);
	def("velocity_cartesian_to_spherical", velocity_cartesian_to_spherical);
	def("add_to_matrix_indexed", add_to_matrix_indexed);
}

void add_to_matrix_indexed(double_matrix M, int_vector _i, int_vector _j, double_vector weight) {
	double* Mp = M.data().begin();
	double* wp = weight.data().begin();
	int* ip = _i.data().begin();
	int* jp = _j.data().begin();
	int size1 = M.size1();
	int size2 = M.size2();
	int length = min(min(_i.size(), _j.size()), weight.size());
	for(int k = 0; k < length; k++) {
		int i = *ip++;
		int j = *jp++;
		double w = *wp++;
		if((i < size1) && (j < size2) && (i >= 0) && (j >= 0)) {
			int index = j + i * size2;
			Mp[index] += w;
		}
	}
}

void matmul_2d_3d(double_matrix M, double_vector x, double_vector y, double_vector xt, double_vector yt, double_vector zt) {
	double* Mp = M.data().begin();
	int Mstride = M.size2();
	double* xp = x.data().begin();
	double* yp = y.data().begin();
	double* xtp = xt.data().begin();
	double* ytp = yt.data().begin();
	double* ztp = zt.data().begin();
	int size = x.size();
	for(int i = 0; i < size; i++) {
		//for(int j = 0; j < 3; j++) {
		*xtp++ = (*xp) * Mp[0+Mstride*0] + (*yp) * Mp[0+Mstride*1];
		*ytp++ = (*xp) * Mp[1+Mstride*0] + (*yp) * Mp[1+Mstride*1];
		*ztp++ = (*xp) * Mp[2+Mstride*0] + (*yp) * Mp[2+Mstride*1];
		xp++; yp++;
		//}
	}
}
void matmul_3d_3d(double_matrix M, double_vector x, double_vector y, double_vector z, double_vector xt, double_vector yt, double_vector zt) {
	double* Mp = M.data().begin();
	int Mstride = M.size2();
	double* xp = x.data().begin();
	double* yp = y.data().begin();
	double* zp = z.data().begin();
	double* xtp = xt.data().begin();
	double* ytp = yt.data().begin();
	double* ztp = zt.data().begin();
	int size = x.size();
	for(int i = 0; i < size; i++) {
		*xtp++ = (*xp) * Mp[0+Mstride*0] + (*yp) * Mp[0+Mstride*1] + (*zp) * Mp[0+Mstride*2];
		*ytp++ = (*xp) * Mp[1+Mstride*0] + (*yp) * Mp[1+Mstride*1] + (*zp) * Mp[1+Mstride*2];
		*ztp++ = (*xp) * Mp[2+Mstride*0] + (*yp) * Mp[2+Mstride*1] + (*zp) * Mp[2+Mstride*2];
		xp++; yp++; zp++;
	}
}

void velocity_cartesian_to_spherical(double_vector xv, double_vector yv, double_vector zv, double_vector vxv, double_vector vyv, double_vector vzv, double_vector vrv, double_vector vphiv, double_vector vthetav) {
	double *xp = xv.data().begin();
	double *yp = yv.data().begin();
	double *zp = zv.data().begin();
	double *vxp = vxv.data().begin();
	double *vyp = vyv.data().begin();
	double *vzp = vzv.data().begin();
	double *vrp = vrv.data().begin();
	double *vphip = vphiv.data().begin();
	double *vthetap = vthetav.data().begin();
	double *xpend = xv.data().end();
	while(xp != xpend) {
		double x = *xp++;
		double y = *yp++;
		double z = *zp++;
		double vx = *vxp++;
		double vy = *vyp++;
		double vz = *vzp++;
		double r = sqrt(x*x+y*y+z*z);
		double rhosq = (x*x+y*y);
		double rho = sqrt(rhosq);
		//double cosatan2xy = x / rho;
		//double sinatan2xy = y / rho;
		*vrp++ = (vx*x + vy*y + vz*z)/r;
		*vphip++ = (vy*x - vx*y)/rho;
		*vthetap++ = (vx*z*x/rho + vy*z*y/rho - vz*rho)/r;

	}
}

}


#include <iostream>

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/timer.hpp>

#include <vector>
#include <math.h>

#include "NRGclasses.hpp"
#include "NRGfunctions.hpp"
#include "TwoChQSz.hpp"

void TwoChQSz_UpdateQm1fQ(CNRGbasisarray* pSingleSite,CNRGarray* pAeig,
			  CNRGbasisarray* pAbasis,CNRGmatrix* Qm1fNQ){

  // Boost matrices
  boost::numeric::ublas::matrix<double> Zibl;
  boost::numeric::ublas::matrix<double> Zjbl;
  boost::numeric::ublas::matrix<double> fnbasis;

  boost::numeric::ublas::matrix<double> fnw;
  // Check time
  boost::timer t;
  double time_elapsed;

  //Clear all
  for (int ich=1;ich<=2;ich++)
    {
      Qm1fNQ[ich-1].ClearAll();
      Qm1fNQ[ich-1].SyncNRGarray(*pAeig);
    }

  // Note: assumes  Abasis and Aeig have
  // EXACTLY the same block structure... should check.

  double Qi[2];
  double Qj[2];
  double Szi,Szj;

  // Loop over blocks (icount counts for each matrix)
  int icount[2]={0,0};
  for (int ibl=0;ibl<pAeig->NumBlocks();ibl++)
    {
      Qi[0]=pAeig->GetQNumber(ibl,0);
      Qi[1]=pAeig->GetQNumber(ibl,1);
      Szi=pAeig->GetQNumber(ibl,2);
      int Nstibl=pAeig->GetBlockSize(ibl);

      for (int jbl=ibl+1;jbl<pAeig->NumBlocks();jbl++)
	{
	  Qj[0]=pAeig->GetQNumber(jbl,0);
	  Qj[1]=pAeig->GetQNumber(jbl,1);
	  Szj=pAeig->GetQNumber(jbl,2);
	  int Nstjbl=pAeig->GetBlockSize(jbl);

	  //Loop over channels: check if the matrix elements are non-zero
	  for (int ich=1;ich<=2;ich++)
	    {
	      if (  dEqual(Qi[ich-1]+1.0,Qj[ich-1])&&
		   ( dEqual(Szi+0.5,Szj)||dEqual(Szi-0.5,Szj) )  )
		{
		  // Set Zi,Zj
		  cout << " BC: Nshell = " << pAeig->Nshell << endl;
		  cout << "Matrix in channel: " << ich << endl;
		  cout << "  Setting up Block i : " << ibl
		       << " (size " <<  Nstibl
		       << ") x Block j " << jbl
		       << " (size " << Nstjbl << ") of " << pAeig->NumBlocks() << endl;
		  Qm1fNQ[ich-1].MatBlockMap.push_back(ibl);
		  Qm1fNQ[ich-1].MatBlockMap.push_back(jbl);
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]);
		  icount[ich-1]+=Nstibl*Nstjbl;
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]-1);

		  cout << "    Setting up BLAS matrices..." << endl;
		  //boost::numeric::ublas::matrix<double> Zibl(Nstibl,Nstibl);
		  //boost::numeric::ublas::matrix<double> Zjbl(Nstjbl,Nstjbl);
		  Zibl.resize(Nstibl,Nstibl);
		  Zjbl.resize(Nstjbl,Nstjbl);

		  Zibl=pAeig->EigVec2BLAS(ibl);
		  Zjbl=pAeig->EigVec2BLAS(jbl);
		  cout << "    ...Zs done ..." << endl;

		  //boost::numeric::ublas::matrix<double> fnbasis (Nstibl,Nstjbl);
		  fnbasis.resize(Nstibl,Nstjbl);
		  // Set-up fn basis: Loop in the basis states
		  int istbl=0;
		  for (int ist=pAbasis->GetBlockLimit(ibl,0);
		       ist<=pAbasis->GetBlockLimit(ibl,1);ist++)
		    {
		      int typei=pAbasis->iType[ist];
		      int stcfi=pAbasis->StCameFrom[ist];

		      int jstbl=0;
		      for (int jst=pAbasis->GetBlockLimit(jbl,0);
			   jst<=pAbasis->GetBlockLimit(jbl,1);jst++)
			{
			  int typej=pAbasis->iType[jst];
			  int stcfj=pAbasis->StCameFrom[jst];

			  fnbasis(istbl,jstbl)=0.0;

			  if (stcfi==stcfj)
			    {
			    fnbasis(istbl,jstbl)=fd_table(ich,-1,typej,typei)+
			      fd_table(ich,1,typej,typei);
			    }
			  jstbl++;
			}
		      // end jst loop
		      istbl++;
		    }
		  // end ist loop
		  cout << "    ...fnbasis done." << endl;

		  cout << "    Multiplying BLAS matrices... " << endl;

		  fnw.resize(Nstibl,Nstjbl);
		  t.restart();
		  noalias(fnw)=prod (Zibl,
				     boost::numeric::ublas::matrix<double>(prod(fnbasis,trans(Zjbl))) );
		  time_elapsed=t.elapsed();
		  cout << "    ...done. Elapsed time:" << time_elapsed << endl;

// 		  cout << "Z(i="<<ibl<<")  : " <<  Zibl << endl;
// 		  cout << "Z(j="<<jbl<<")  : " <<  Zjbl << endl;
// 		  cout << "fbasis   :" <<  fnbasis << endl;
// 		  cout << "Zi.fbasis.ZjT : " <<  fnw << endl;

		  // Add to Qm1fNQ[ich-1]
		  for (int ii=0;ii<fnw.size1();ii++)
		    for (int jj=0;jj<fnw.size2();jj++)
		      Qm1fNQ[ich-1].MatEl.push_back(fnw(ii,jj));

		}
	      //end if Q=Q'+1 etc

	    }
	  // end channel loop

	}
      // end jbl loop

    }
  // end ibl loop

}
// END subroutine

//////////////////////////////////
//////////////////////////////////
//////////////////////////////////
//////////////////////////////////

void TwoChQSz_UpdateMatrixAfterCutting(CNRGbasisarray* pSingleSite,
				       CNRGbasisarray* pAeigCut,
				       CNRGbasisarray* pAbasis,
				       CNRGmatrix* Qm1fNQ,
				       CNRGmatrix* pMQQp1){

  // Boost matrices
  boost::numeric::ublas::matrix<double> Zibl;
  boost::numeric::ublas::matrix<double> Zjbl;
  boost::numeric::ublas::matrix<double> fnbasis;

  boost::numeric::ublas::matrix<double> fnw;
  // Check time
  boost::timer t;
  double time_elapsed;

  //Clear all
  for (int ich=1;ich<=2;ich++)
    {
      Qm1fNQ[ich-1].ClearAll();
      Qm1fNQ[ich-1].SyncNRGarray(*pAeigCut);
    }

  // Note: assumes  Abasis and Aeig have
  // EXACTLY the same block structure... should check.

  double Qi[2];
  double Qj[2];
  double Szi,Szj;

  double qnums[3];

  // Loop over blocks (icount counts for each matrix)
  int icount[2]={0,0};
  for (int ibl=0;ibl<pAeigCut->NumBlocks();ibl++)
    {
      Qi[0]=pAeigCut->GetQNumber(ibl,0);
      Qi[1]=pAeigCut->GetQNumber(ibl,1);
      Szi=pAeigCut->GetQNumber(ibl,2);
      int Nstibl=pAeigCut->GetBlockSize(ibl);

      // Get the corresponding block in pAbasis!
      qnums[0]=Qi[0];
      qnums[1]=Qi[1];
      qnums[2]=Szi;
      int ii_corr=pAbasis->GetBlockFromQNumbers(qnums);
      int NstiblBC=pAbasis->GetBlockSize(ii_corr);

      for (int jbl=ibl+1;jbl<pAeigCut->NumBlocks();jbl++)
	{
	  Qj[0]=pAeigCut->GetQNumber(jbl,0);
	  Qj[1]=pAeigCut->GetQNumber(jbl,1);
	  Szj=pAeigCut->GetQNumber(jbl,2);
	  int Nstjbl=pAeigCut->GetBlockSize(jbl);

	  // Get the corresponding block in pAbasis!
	  qnums[0]=Qj[0];
	  qnums[1]=Qj[1];
	  qnums[2]=Szj;
	  int jj_corr=pAbasis->GetBlockFromQNumbers(qnums);
	  int NstjblBC=pAbasis->GetBlockSize(jj_corr);

	  if ( (!dEqual(Qj[0],pAbasis->GetQNumber(jj_corr,0)))||
	       (!dEqual(Qj[1],pAbasis->GetQNumber(jj_corr,1)))||
	       (!dEqual(Szj,pAbasis->GetQNumber(jj_corr,2))) )
	    {
	      cout << "Ops, problems with GetBlockFromQNumbers" << endl;
	      pAbasis->PrintQNumbers();
	      cout << qnums[0] << " " << qnums[1] << " " << qnums[2] << endl;
	      cout << pAbasis->GetBlockFromQNumbers(qnums) << endl;
	    }

	  //Loop over channels: check if the matrix elements are non-zero
	  for (int ich=1;ich<=2;ich++)
	    {
	      if (  dEqual(Qi[ich-1]+1.0,Qj[ich-1])&&
		   ( dEqual(Szi+0.5,Szj)||dEqual(Szi-0.5,Szj) )  )
		{
		  // Set Zi,Zj
		  cout << " AC : Nshell = " << pAeigCut->Nshell << endl;
		  cout << "Matrix in channel: " << ich << endl;
		  cout << "  Setting up Block i : " << ibl
		       << " (size " <<  Nstibl << " was " << NstiblBC
		       << ") x Block j " << jbl
		       << " (size " << Nstjbl << " was " << NstjblBC
		       << ") of " << pAeigCut->NumBlocks() << endl;
		  Qm1fNQ[ich-1].MatBlockMap.push_back(ibl);
		  Qm1fNQ[ich-1].MatBlockMap.push_back(jbl);
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]);
		  icount[ich-1]+=Nstibl*Nstjbl;
		  Qm1fNQ[ich-1].MatBlockBegEnd.push_back(icount[ich-1]-1);

		  cout << "    Setting up BLAS matrices..." << endl;
		  //boost::numeric::ublas::matrix<double> Zibl(Nstibl,Nstibl);
		  //boost::numeric::ublas::matrix<double> Zjbl(Nstjbl,Nstjbl);
		  Zibl.resize(Nstibl,NstiblBC);
		  Zjbl.resize(Nstjbl,NstjblBC);

		  //Zibl=pAeigCut->EigVec2BLAS(ibl);
		  //Zjbl=pAeigCut->EigVec2BLAS(jbl);
		  Zibl=pAeigCut->EigVecCut2BLAS(ibl);
		  Zjbl=pAeigCut->EigVecCut2BLAS(jbl);
		  cout << "    ...Zs done ..." << endl;

		  fnbasis.resize(NstiblBC,NstjblBC);
		  // Set-up fn basis: Loop in the basis states
		  // Watch out here... ii_corr, not ibl
		  int istbl=0;
		  for (int ist=pAbasis->GetBlockLimit(ii_corr,0);
		       ist<=pAbasis->GetBlockLimit(ii_corr,1);ist++)
		    {
		      int typei=pAbasis->iType[ist];
		      int stcfi=pAbasis->StCameFrom[ist];

		      int jstbl=0;
		      for (int jst=pAbasis->GetBlockLimit(jj_corr,0);
			   jst<=pAbasis->GetBlockLimit(jj_corr,1);jst++)
			{
			  int typej=pAbasis->iType[jst];
			  int stcfj=pAbasis->StCameFrom[jst];

			  fnbasis(istbl,jstbl)=0.0;

			  if (stcfi==stcfj)
			    {
			    fnbasis(istbl,jstbl)=fd_table(ich,-1,typej,typei)+
			      fd_table(ich,1,typej,typei);
			    }
			  jstbl++;
			}
		      // end jst loop
		      istbl++;
		    }
		  // end ist loop
		  cout << "    ...fnbasis done." << endl;

		  cout << "    Multiplying BLAS matrices... " << endl;

		  fnw.resize(Nstibl,Nstjbl);
		  t.restart();
		  noalias(fnw)=prod (Zibl,
				     boost::numeric::ublas::matrix<double>(prod(fnbasis,trans(Zjbl))) );
		  time_elapsed=t.elapsed();
		  cout << "    ...done. Elapsed time:" << time_elapsed << endl;

// 		  cout << "Z(i="<<ibl<<")  : " <<  Zibl << endl;
// 		  cout << "Z(j="<<jbl<<")  : " <<  Zjbl << endl;
// 		  cout << "fbasis   :" <<  fnbasis << endl;
// 		  cout << "Zi.fbasis.ZjT : " <<  fnw << endl;

		  // Add to Qm1fNQ[ich-1]
		  for (int ii=0;ii<fnw.size1();ii++)
		    for (int jj=0;jj<fnw.size2();jj++)
		      Qm1fNQ[ich-1].MatEl.push_back(fnw(ii,jj));

		}
	      //end if Q=Q'+1 etc

	    }
	  // end channel loop

	}
      // end jbl loop

    }
  // end ibl loop

}
//////////////////////////////////
//////////////////////////////////
//////////////////////////////////
//////////////////////////////////


// =========================================================================
// @author Leonardo Florez-Valencia (florez-l@javeriana.edu.co)
// =========================================================================

#include "NeuralNetwork.h"

#include <cassert>
#include <map>
#include <regex>
#include <sstream>
#include <boost/algorithm/string.hpp>

// -------------------------------------------------------------------------
template< class _TScl >
NeuralNetwork< _TScl >::
NeuralNetwork( )
{
}

// -------------------------------------------------------------------------
template< class _TScl >
NeuralNetwork< _TScl >::
NeuralNetwork( const Self& o )
{
  this->m_L.clear( );
  this->m_L.insert( this->m_L.begin( ), o.m_L.begin( ), o.m_L.end( ) );
}

// -------------------------------------------------------------------------
template< class _TScl >
typename NeuralNetwork< _TScl >::
Self& NeuralNetwork< _TScl >::
operator=( const Self& o )
{
  this->m_L.clear( );
  this->m_L.insert( this->m_L.begin( ), o.m_L.begin( ), o.m_L.end( ) );
  return( *this );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
add( unsigned int i, unsigned int o, const std::string& f )
{
  this->add( TLayer( i, o, f ) );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
add( unsigned int o, const std::string& f )
{
  assert( this->m_L.size( ) > 0 && "At least one layer is needed" );

  this->add( this->m_L.back( ).output_size( ), o, f );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
add( const TMatrix& w, const TColVector& b, const std::string& f )
{
  this->add( TLayer( w, b, f ) );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
add( const TLayer& l )
{
  if( this->m_L.size( ) > 0 )
    assert(
      l.input_size( ) == this->m_L.back( ).output_size( ) && "Invalid sizes"
      );
  this->m_L.push_back( l );

  // Normalization transforms
  if( this->m_L.size( ) == 1 )
  {
    unsigned int n = this->m_L[ 0 ].input_size( );
    this->m_NormalizationOffset = TColVector::Zero( n );
    this->m_NormalizationScale = TMatrix::Identity( n, n );
  } // end if
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
load_topology( std::istream& is )
{
  unsigned int L;
  std::map< unsigned int, unsigned long > K;
  std::map< unsigned int, std::string > F;

  // Read parameters
  std::regex r( "\\s+" );
  std::string line;
  while( std::getline( is, line ) )
  {
    line = std::regex_replace( line, r, "" );
    std::vector< std::string > tokens;
    boost::algorithm::split( tokens, line, boost::is_any_of( "=" ) );
    if( tokens.size( ) == 2 )
    {
      if( tokens[ 0 ] == "L" )
      {
        std::istringstream data( tokens[ 1 ] );
        data >> L;
      }
      else
      {
        if( tokens[ 0 ][ 0 ] == 'k' )
        {
          std::istringstream i_str( tokens[ 0 ].substr( 1 ) );
          std::istringstream k_str( tokens[ 1 ] );
          unsigned int i;
          unsigned long k;
          i_str >> i;
          k_str >> k;
          K[ i ] = k;
        }
        else if( tokens[ 0 ][ 0 ] == 'f' )
        {
          std::istringstream i_str( tokens[ 0 ].substr( 1 ) );
          unsigned int i;
          i_str >> i;
          F[ i ] = tokens[ 1 ];
        } // end if
      } // end if
    } // end if
  } // end while

  // Real build
  if( F.size( ) == L && K.size( ) == L + 1 )
    for( unsigned int l = 0; l < L; ++l )
      this->add( K[ l ], K[ l + 1 ], F[ l ] );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
set( unsigned int l, const TMatrix& w, const TColVector& b )
{
  if( l < this->m_L.size( ) )
  {
    this->m_L[ l ].weights( ) = w;
    this->m_L[ l ].biases( ) = b;
  } // end if
}

// -------------------------------------------------------------------------
template< class _TScl >
unsigned int NeuralNetwork< _TScl >::
number_of_layers( ) const
{
  return( this->m_L.size( ) );
}

// -------------------------------------------------------------------------
template< class _TScl >
typename NeuralNetwork< _TScl >::
TMatrix& NeuralNetwork< _TScl >::
weights( unsigned int l )
{
  static TMatrix w( 1, 1 );
  if( l < this->m_L.size( ) )
    return( this->m_L[ l ].weights( ) );
  else
    return( w );
}

// -------------------------------------------------------------------------
template< class _TScl >
const typename NeuralNetwork< _TScl >::
TMatrix& NeuralNetwork< _TScl >::
weights( unsigned int l ) const
{
  static const TMatrix w( 1, 1 );
  if( l < this->m_L.size( ) )
    return( this->m_L[ l ].weights( ) );
  else
    return( w );
}

// -------------------------------------------------------------------------
template< class _TScl >
typename NeuralNetwork< _TScl >::
TColVector& NeuralNetwork< _TScl >::
biases( unsigned int l )
{
  static TColVector b( 1 );
  if( l < this->m_L.size( ) )
    return( this->m_L[ l ].biases( ) );
  else
    return( b );
}

// -------------------------------------------------------------------------
template< class _TScl >
const typename NeuralNetwork< _TScl >::
TColVector& NeuralNetwork< _TScl >::
biases( unsigned int l ) const
{
  static const TColVector b( 1 );
  if( l < this->m_L.size( ) )
    return( this->m_L[ l ].biases( ) );
  else
    return( b );
}

// -------------------------------------------------------------------------
template< class _TScl >
typename NeuralNetwork< _TScl >::
TActivation* NeuralNetwork< _TScl >::
sigma( unsigned int l )
{
  if( l < this->m_L.size( ) )
    return( this->m_L[ l ].sigma( ) );
  else
    return( nullptr );
}

// -------------------------------------------------------------------------
template< class _TScl >
const typename NeuralNetwork< _TScl >::
TActivation* NeuralNetwork< _TScl >::
sigma( unsigned int l ) const
{
  if( l < this->m_L.size( ) )
    return( this->m_L[ l ].sigma( ) );
  else
    return( nullptr );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
setNormalizationOffset( const TColVector& o )
{
  this->m_NormalizationOffset = o;
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
setNormalizationScale( const TMatrix& s )
{
  this->m_NormalizationScale = s;
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
init( )
{
  // Weights
  for( TLayer& l: this->m_L )
    l.init( );

  // Normalization transforms
  unsigned int n = this->m_L[ 0 ].input_size( );
  this->m_NormalizationOffset = TColVector::Zero( n );
  this->m_NormalizationScale = TMatrix::Identity( n, n );
}

// -------------------------------------------------------------------------
template< class _TScl >
typename NeuralNetwork< _TScl >::
TMatrix NeuralNetwork< _TScl >::
f( const TMatrix& x ) const
{
  typename TLayers::const_iterator lIt = this->m_L.begin( );
  TMatrix z = lIt->f(
    this->m_NormalizationScale *
    ( x.colwise( ) + this->m_NormalizationOffset )
    );
  for( lIt++; lIt != this->m_L.end( ); ++lIt )
    z = lIt->f( z );
  return( z );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
f( std::vector< TMatrix >& a, std::vector< TMatrix >& z ) const
{
  typename TLayers::const_iterator lIt = this->m_L.begin( );
  typename std::vector< TMatrix >::iterator aIt, bIt, zIt;
  aIt = bIt = a.begin( );
  zIt = z.begin( );

  for( bIt++; lIt != this->m_L.end( ); ++lIt, ++aIt, ++bIt, ++zIt )
    *bIt = lIt->f( *aIt, *zIt );
}

// -------------------------------------------------------------------------
template< class _TScl >
typename NeuralNetwork< _TScl >::
TMatrix NeuralNetwork< _TScl >::
t( const TMatrix& x ) const
{
  return( this->m_L.back( ).sigma( )->t( this->f( x ) ) );
}

// -------------------------------------------------------------------------
template< class _TScl >
typename NeuralNetwork< _TScl >::
TMatrix NeuralNetwork< _TScl >::
_d( const unsigned int& l, const TMatrix& z ) const
{
  assert( l < this->m_L.size( ) && "Layer does not exist" );

  return( this->m_L[ l ].sigma( )->d( z ) );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
_read_from( std::istream& i )
{
  unsigned int L;
  i >> L;
  for( unsigned int n = 0; n < L; ++n )
  {
    TLayer l;
    i >> l;
    this->add( l );
  } // end for

  // Read normalization parameters
  unsigned int N = this->m_L[ 0 ].input_size( );
  this->m_NormalizationOffset = TColVector::Zero( N );
  this->m_NormalizationScale = TMatrix::Identity( N, N );
  for( unsigned int n = 0; n < N; ++n )
    i >> this->m_NormalizationOffset( n, 0 );
  for( unsigned int x = 0; x < N; ++x )
    for( unsigned int y = 0; y < N; ++y )
      i >> this->m_NormalizationScale( x, y );
}

// -------------------------------------------------------------------------
template< class _TScl >
void NeuralNetwork< _TScl >::
_copy_to( std::ostream& o ) const
{
  o << this->m_L.size( ) << std::endl;
  for( const TLayer& l: this->m_L )
    o << l << std::endl;
  o << this->m_NormalizationOffset << std::endl;
  o << this->m_NormalizationScale << std::endl;
}

// -------------------------------------------------------------------------
template class NeuralNetwork< float >;
template class NeuralNetwork< double >;
template class NeuralNetwork< long double >;

// eof - $RCSfile$


#include "tkdCmdParser.h"

#include "vnl/vnl_matrix.h"
#include "vnl/vnl_vector.h"
#include "vnl/algo/vnl_determinant.h"
#include "vnl/algo/vnl_svd.h"

#include <stdlib.h>
#include <math.h>
#include <fstream>

#include <boost/tokenizer.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/uniform_real.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/random.hpp>
#include <boost/utility.hpp>
#include <boost/accumulators/numeric/functional/vector.hpp>
#include <boost/accumulators/numeric/functional/complex.hpp>
#include <boost/accumulators/numeric/functional/valarray.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/variance.hpp>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics.hpp>

#include <boost/math/distributions/normal.hpp>

/* Indexer problem... */

/**
 * Option list.
 */
struct parameters
{
	std::string inputFileName;
	std::string link;
	std::string separation;
	int iterations;
	double percentage;
	bool verbose;
	bool printPValues;
	bool printZValues;
};

/**
 * Generalized Linear Model (GLM) is a flexible generalization of ordinary least squares regression.
 *
 * http://en.wikipedia.org/wiki/Generalized_linear_model
 *
 * wim@invivonmr.uu.nl, 02-03-2010.
 *
 */
namespace generalized_linear_model
{
	typedef double PixelType;
	typedef vnl_matrix< PixelType > MatrixType;
	typedef vnl_vector< PixelType > VectorType;
	typedef boost::tokenizer< boost::char_separator< char > > TokType;
	typedef boost::mt19937 random_number_type;
	typedef boost::uniform_int< long > int_distribution_type;
	typedef boost::variate_generator< random_number_type&, int_distribution_type > int_generator_type;

	class GLM
	{
	public:

		/**
		 * Run glm.
		 */
		void Run( parameters& list )
		{
			// read covariates and response-function (latter as last column).
			MatrixType covariates;
			VectorType binaryResponse;

			if( list.verbose )
				std::cout << "Reading input data..." << std::endl;

			ReadData( list.separation, list.inputFileName, covariates, binaryResponse, list.percentage );
			covariates = AddIntercept( covariates );

			if( list.verbose )
				std::cout << "Finished reading input data..." << std::endl;

			if ( list.link == "logistic" )
			{
				// Calculate coefficients...
				VectorType coefficients = IRLS( binaryResponse, covariates, list.iterations, list.verbose );
				std::cout << "Coefficients: " << coefficients << std::endl;

				if( list.printZValues )
					std::cout << "Z-values: " << ZStatistic( coefficients, binaryResponse, covariates ) << std::endl;

				if( list.printPValues )
					std::cout << "P-values: " << PValue( ZStatistic( coefficients, binaryResponse, covariates ) ) << std::endl;
			}
			else
			{
				std::cerr << "*** ERROR ***: Link function not supported!" << std::endl;
				exit( EXIT_FAILURE );
			}
		}

	protected:

		/**
		 * The P values for individual coefficients, given the z values.
		 */
		VectorType PValue( const VectorType& zValues )
		{
			VectorType pValues( zValues.size(), 0.0 );

			for ( unsigned int i = 0; i < zValues.size(); i++ )
			{
				// two-sided
				boost::math::normal s;
				PixelType alpha = boost::math::pdf( s, zValues( i ) ) * 2.;
				pValues( i ) = alpha;
			}
			return pValues;
		}

		/**
		 * The z statistics for individual coefficients.
		 */
		VectorType ZStatistic( const VectorType& coefficients, const VectorType& response,
				const MatrixType& covariates )
		{
			VectorType means = Means( coefficients, covariates );
			MatrixType weights = Weights( means, coefficients, covariates );

			MatrixType xwx = covariates * weights * covariates.transpose();

			vnl_svd< double > svd( xwx );

			MatrixType variance = svd.inverse();

			VectorType coefficientSE( variance.rows(), 0.0 );

			for( unsigned int i = 0; i < variance.rows(); i++ )
				coefficientSE( i ) = std::sqrt( variance( i, i ) );

			MatrixType correlation( variance.rows(), variance.rows(), 0.0 );

			for( unsigned int i = 0; i < variance.rows(); i++)
			{
				for ( unsigned int j = i; j < variance.rows(); j++ )
				{
					correlation( i, j ) = variance( i, j ) /
											std::sqrt(variance( i, i ) * variance( j, j ) );

					correlation( j, i ) = correlation( i, j );
				}
			}

			VectorType zValues( coefficients.size(), 0.0 );

			for( unsigned int i = 0; i < zValues.size(); i++ )
				zValues( i ) = coefficients( i ) / coefficientSE( i );

			return zValues;
		}

		/**
		 * Return number of training dimensions.
		 */
		unsigned int GetNumberOfTrainingDims( const std::string& trainingData )
		{
			std::string sep = "\t";
			std::ifstream in( trainingData.c_str() );

			if ( in.fail() )
			{
				std::cerr << "*** ERROR ***: could not read labels from: " << trainingData << std::endl;
				exit( EXIT_FAILURE );
			}

			std::string line;
			getline( in, line );

			TokType tok( line, boost::char_separator< char >( sep.c_str() ) );

			unsigned int results = 0;
			try
			{
				for ( TokType::iterator id = tok.begin(); id != tok.end(); ++id )
					results++;
			}
			catch ( boost::bad_lexical_cast& e )
			{
				std::cout << "*** WARNING ***: bad lexical cast during training data parsing!" << std::endl;
			}
			in.close();
			return results;
		}

		/**
		 * Return number of training data points.
		 */
		unsigned int GetNumberOfTrainingPoints( const std::string& trainingData )
		{
			std::ifstream in( trainingData.c_str() );

			if ( in.fail() )
			{
				std::cerr << "*** ERROR ***: could not read labels from: " << trainingData << std::endl;
				exit( EXIT_FAILURE );
			}
			std::string line;
			unsigned int result = 0;

			while ( getline( in, line ) )
			{
				result++;
			}

			in.close();
			return result;
		}

		/**
		 * Add column with initial intercept values (1.0).
		 */
		MatrixType AddIntercept( const MatrixType& M )
		{
			MatrixType result( M.rows() + 1, M.cols() );

			for( unsigned int i = 0; i < result.rows(); i++ )
			{
				for( unsigned int j = 0; j < result.cols(); j++ )
				{
					if( i == 0 )
						result( i, j ) = 1.0;
					else
						result( i, j ) = M( i - 1, j );
				}
			}
			return result;
		}

		/**
		 * Read data from text file.
		 */
		void ReadData( const std::string& sep, const std::string& inputFileName, MatrixType& covariates, VectorType& binaryResponse, PixelType percentage )
		{
			unsigned int dims = GetNumberOfTrainingDims( inputFileName );
			unsigned int points = GetNumberOfTrainingPoints( inputFileName );

			covariates = MatrixType( dims - 1, points, 0.0 );
			binaryResponse = VectorType( points, 0.0 );

			// open training data ...
			std::ifstream in( inputFileName.c_str() );
			if ( in.fail() )
			{
				std::cerr << "*** ERROR ***: could not read data from: " << inputFileName << std::endl;
				exit( EXIT_FAILURE );
			}

			// insert all values into matrix, except the last one (responses) which is inserted in
			// the respose vector ...
			std::string line;
			unsigned int rowIndex = 0;
			while ( getline( in, line ) )
			{
				TokType tok( line, boost::char_separator< char >( sep.c_str() ) );
				try
				{
					unsigned int colIndex = 0;
					for ( TokType::iterator id = tok.begin(); id != tok.end(); ++id )
					{
						if ( colIndex == dims - 1 ) // last column ...
							binaryResponse( rowIndex ) = boost::lexical_cast< PixelType >( *id );
						else
							covariates( colIndex, rowIndex ) = boost::lexical_cast< PixelType >( *id );

						colIndex++;
					}
					rowIndex++;
				} catch ( boost::bad_lexical_cast& e )
				{
					std::cout << "*** WARNING ***: could not parse data! (is it in column format?)" << std::endl;
					e.what();
					exit( EXIT_FAILURE );
				}
			}
			in.close();

			if ( ( percentage > 0 ) && ( percentage < 100 ) )
			{
				ResizeTrainingData( covariates, binaryResponse, percentage );
			}
		}

		/**
		 * Reduce number of training-points to given percentage.
		 */

		void ResizeTrainingData( MatrixType& trainingData, VectorType& trainingClasses,
				PixelType percentage )
		{
			unsigned int pointsToKeep = ( trainingClasses.size() * percentage ) / 100.0;

			// initialize vector ...
			VectorType subClasses( pointsToKeep, 0.0 );

			// initialize matrix ...
			MatrixType subData( trainingData.rows(), pointsToKeep, 0.0 );

			// create generator ...
			random_number_type generator( time( 0 ) );

			VectorType indices = GetRandomValues( trainingClasses, generator, pointsToKeep );

			// insert random samples ...
			for ( unsigned int i = 0; i < indices.size(); i++ )
			{
				subClasses( i ) = trainingClasses( indices( i ) );

				for ( unsigned int j = 0; j < subData.rows(); j++ )
					subData( j, i ) = trainingData( j, indices( i ) );
			}

			trainingData = subData;
			trainingClasses = subClasses;
		}

		/**
		 * Return uniform random values vector from inputs, with similar size.
		 */
		VectorType GetRandomValues( const VectorType& inputs, random_number_type& generator, unsigned int total )
		{
			// uniform distribution: 0 -> largest index ...
			int_distribution_type int_uni_dist( 0, inputs.size() - 1 );

			// generator ...
			int_generator_type int_distribution( generator, int_uni_dist );

			VectorType results( total, 0.0 );

			for ( unsigned int i = 0; i < total; i++ )
				results( int_distribution() );

			return results;
		}

		/**
		 * Return weight matrix.
		 */
		MatrixType Weights( const VectorType& means, const VectorType& coefficients, const MatrixType& covariate )
		{
			MatrixType weights( means.size(), means.size(), 0 );

			for ( unsigned int i = 0; i < means.size(); i++ )
				weights( i, i ) = means( i ) * ( 1. - ( means( i ) ) );

			return weights;
		}

		/**
		 * Return means vector.
		 */
		VectorType Means( const VectorType& coefficients, const MatrixType& covariate )
		{
			VectorType linearPredictors = GetColumnPackedCopy( covariate.transpose() * Vector2Matrix( coefficients ) );

			VectorType means( linearPredictors.size() );

			for ( unsigned int i = 0; i < means.size(); i++ )
				means( i) = std::exp( linearPredictors( i ) ) / ( static_cast< PixelType > ( 1.0 ) + std::exp( linearPredictors( i ) ) );

			return means;
		}

		/**
		 * Construct matrix from vector.
		 */
		MatrixType Vector2Matrix( const VectorType& V )
		{
			MatrixType A( V.size(), 1 );
			A.set_column( 0, V );

			return A;
		}

		/**
		 * Convert all matrix columns to vector.
		 */
		VectorType GetColumnPackedCopy( const MatrixType& M )
		{
			VectorType V( M.rows() * M.cols() );

			for ( unsigned int i = 0; i < M.rows(); i++ )
				for ( unsigned int j = 0; j < M.cols(); j++ )
					V( i + j * M.rows() ) = M( i, j );

			return V;
		}

		/**
		 * Set initial values for IRLS estimation.
		 */
		VectorType SetInit( const MatrixType& covariate )
		{
			VectorType V( covariate.rows() );

			PixelType a = 0;

			for ( unsigned int i = 0; i < V.size(); i++ )
			{
				a = 0;
				for ( unsigned int j = 0; j < covariate.get_row( i ).size(); j++ )
				{
					a += covariate( i, j );
				}

				V( i) = covariate.get_row( i ).size() / ( static_cast< PixelType > ( 100 ) * a );
			}
			return V;
		}

		/**
		 * Reweighted Least Squares ( IRLS ) estimation, by finding the
		 * maximum likelihood estimates of a generalized linear model( GLM ).
		 */
		VectorType IRLS( const VectorType& response, const MatrixType& covariate, unsigned int max_iter, bool verbose )
		{
			if ( response.size() != covariate.cols() )
			{
				std::cerr << "*** ERROR ***: The response vector and rows of the "
					"covariate matrix must have the same length." << std::endl;
				exit( EXIT_FAILURE );
			}

			VectorType coefficients = SetInit( covariate );

			MatrixType responseMatrix = Vector2Matrix( response );

			PixelType error = 1.0;

			MatrixType covariateMatrix( covariate );

			if( verbose )
				std::cout << "Starting iterative solving..." << std::endl;

			unsigned int iter = 0;
			while ( error > 0.000001 )
			{
				if( iter > max_iter )
				{
					std::cerr << "*** ERROR ***: maximum iteration reached. No convergence."<< std::endl;
					exit( EXIT_FAILURE );
				}

				if( verbose )
					std::cout << "Iteration: " << iter << std::endl;

				MatrixType coefficientMatrix = Vector2Matrix( coefficients );

				VectorType means = Means( coefficients, covariate );
				MatrixType weights = Weights( means, coefficients, covariate );

				MatrixType inversedWeights( weights.rows(), weights.cols() );
				MatrixType weight = MatrixType( weights );

				for ( unsigned int i = 0; i < weights.rows(); i++ )
				{
					if ( weights( i, i ) == 0. )
						inversedWeights( i, i ) = 0.;
					else
						inversedWeights( i, i ) = static_cast< PixelType > ( 1. ) / weights( i, i );
				}

				MatrixType linearPredictor = covariateMatrix.transpose() * coefficientMatrix;
				MatrixType z = linearPredictor + ( inversedWeights ) * ( responseMatrix - ( Vector2Matrix( means ) ) );
				MatrixType xwx = covariateMatrix * ( weight * covariateMatrix.transpose() );

				// Correct for rows with all values close to zero, by adding 0.1 to diagonal ...
				if ( std::abs( vnl_determinant< PixelType > ( xwx ) ) <= 1e-8 )
					for ( unsigned int i = 0; i < xwx.rows(); i++ )
						xwx( i, i ) = xwx( i, i ) + 0.1;

				// inverse...
				vnl_svd< PixelType > svd( xwx );

				MatrixType updatedCoefficient = svd.inverse() * ( ( covariateMatrix * weight ) * z );

				coefficients = GetColumnPackedCopy( updatedCoefficient );
				error = std::pow( ( updatedCoefficient - coefficientMatrix ).frobenius_norm(), 2.0 );

				iter++;
			}

			return coefficients;
		}

	};
} // end namespace generalized_linear_model

/**
 * Main.
 */
int main( int argc, char ** argv )
{
	tkd::CmdParser p( argv[0], "Generalized linear model" );

	parameters list;

	list.iterations = 100000;
	list.link = "logistic";
	list.separation = "\t";
	list.percentage = 100.0;
	list.verbose = false;
	list.printPValues = false;
	list.printZValues = false;

	p.AddArgument( list.inputFileName, "input" ) ->AddAlias( "i" ) ->SetInput( "filename" ) ->SetDescription(
			"Input file: list of covariate(s) and response function (column format)" ) ->SetRequired( true );

	p.AddArgument( list.link, "link" ) ->AddAlias( "l" ) ->SetInput( "string" ) ->SetDescription( "Link function (default: \"logistic\")" );

	p.AddArgument( list.separation, "separation" ) ->AddAlias( "s" ) ->SetInput( "string" ) ->SetDescription( "Data separator string (default: \\t" );

	p.AddArgument( list.printZValues, "print-z-values" ) ->AddAlias( "pz" ) ->SetInput( "bool" ) ->SetDescription( "Print Z-values for coefficients (default: false" );

	p.AddArgument( list.printPValues, "print-p-values" ) ->AddAlias( "pp" ) ->SetInput( "bool" ) ->SetDescription( "Print P-values for coefficients (default: false" );

	p.AddArgument( list.percentage, "percentage" ) ->AddAlias( "p" ) ->SetInput( "float" ) ->SetDescription(
			"Percentage of covariates to select randomly (default: 100)" );

	p.AddArgument( list.verbose, "verbose" ) ->AddAlias( "v" ) ->SetInput( "bool" ) ->SetDescription( "Verbose (default: false" );

	if ( !p.Parse( argc, argv ) )
	{
		p.PrintUsage( std::cout );
		return EXIT_FAILURE;
	}

	generalized_linear_model::GLM glm;

	glm.Run( list );

	return EXIT_SUCCESS;
}


// Distributed under the MIT License.
// See LICENSE.txt for details.

#include "Utilities/Rational.hpp"

#include <boost/functional/hash.hpp>
#include <boost/integer/common_factor_rt.hpp>
#include <ostream>
#include <pup.h>
#include <tuple>

#include "Utilities/ErrorHandling/Assert.hpp"

// IWYU pragma: no_include <boost/cstdint.hpp>

namespace {
// This is needed a lot, so define a shorter name.
std::int64_t to64(const std::int32_t n) { return static_cast<std::int64_t>(n); }

template <typename IntType>
std::tuple<std::int32_t, std::int32_t> reduce(IntType numerator,
                                              IntType denominator) {
  const IntType common_factor = boost::integer::gcd(numerator, denominator);
  numerator /= common_factor;
  denominator /= common_factor;
  if (denominator < 0) {
    numerator = -numerator;
    denominator = -denominator;
  }
  ASSERT(static_cast<std::int32_t>(numerator) == numerator and
         static_cast<std::int32_t>(denominator) == denominator,
         "Rational overflow: " << numerator << "/" << denominator);
  return std::make_tuple(numerator, denominator);
}
}  // namespace

Rational::Rational(const std::int32_t numerator,
                   const std::int32_t denominator) {
  ASSERT(denominator != 0, "Division by zero");
  std::tie(numerator_, denominator_) = reduce(numerator, denominator);
}

double Rational::value() const {
  return static_cast<double>(numerator_) / static_cast<double>(denominator_);
}

Rational Rational::inverse() const {
  // Default construct to avoid the reduce() call.
  ASSERT(*this != 0, "Division by zero");
  Rational ret;
  ret.numerator_ = denominator_;
  ret.denominator_ = numerator_;
  if (ret.denominator_ < 0) {
    ret.numerator_ = -ret.numerator_;
    ret.denominator_ = -ret.denominator_;
  }
  return ret;
}

Rational& Rational::operator+=(const Rational& other) {
  std::tie(numerator_, denominator_) =
      reduce(to64(numerator_) * to64(other.denominator_) +
             to64(denominator_) * to64(other.numerator_),
             to64(denominator_) * to64(other.denominator_));
  return *this;
}
Rational& Rational::operator-=(const Rational& other) {
  return *this += -other;
}
Rational& Rational::operator*=(const Rational& other) {
  std::tie(numerator_, denominator_) =
      reduce(to64(numerator_) * to64(other.numerator()),
             to64(denominator_) * to64(other.denominator()));
  return *this;
}
Rational& Rational::operator/=(const Rational& other) {
  return *this *= other.inverse();
}

void Rational::pup(PUP::er& p) {
  p | numerator_;
  p | denominator_;
}

Rational operator-(Rational r) {
  // No reduced-form check needed
  r.numerator_ = -r.numerator_;
  return r;
}

Rational operator+(const Rational& a, const Rational& b) {
  Rational ret = a;
  ret += b;
  return ret;
}
Rational operator-(const Rational& a, const Rational& b) {
  Rational ret = a;
  ret -= b;
  return ret;
}
Rational operator*(const Rational& a, const Rational& b) {
  Rational ret = a;
  ret *= b;
  return ret;
}
Rational operator/(const Rational& a, const Rational& b) {
  Rational ret = a;
  ret /= b;
  return ret;
}

bool operator==(const Rational& a, const Rational& b) {
  return a.numerator() == b.numerator() and a.denominator() == b.denominator();
}
bool operator!=(const Rational& a, const Rational& b) { return not(a == b); }
bool operator<(const Rational& a, const Rational& b) {
  return to64(a.numerator()) * to64(b.denominator()) <
         to64(b.numerator()) * to64(a.denominator());
}
bool operator>(const Rational& a, const Rational& b) { return b < a; }
bool operator<=(const Rational& a, const Rational& b) { return not(b < a); }
bool operator>=(const Rational& a, const Rational& b) { return not(a < b); }

std::ostream& operator<<(std::ostream& os, const Rational& r) {
  return os << r.numerator() << '/' << r.denominator();
}

size_t hash_value(const Rational& r) {
  size_t h = 0;
  boost::hash_combine(h, r.numerator());
  boost::hash_combine(h, r.denominator());
  return h;
}

// clang-tidy: do not modify std namespace (okay for hash)
namespace std {  // NOLINT
size_t hash<Rational>::operator()(const Rational& r) const {
  return boost::hash<Rational>{}(r);
}
}  // namespace std


/* =========================================================================
   Copyright (c) 2010-2014, Institute for Microelectronics,
                            Institute for Analysis and Scientific Computing,
                            TU Wien.
   Portions of this software are copyright by UChicago Argonne, LLC.

                            -----------------
                  ViennaCL - The Vienna Computing Library
                            -----------------

   Project Head:    Karl Rupp                   rupp@iue.tuwien.ac.at

   (A list of authors and contributors can be found in the PDF manual)

   License:         MIT (X11), see file LICENSE in the base directory
============================================================================= */

/*
*
*   Tutorial: Using ViennaCL with multiple threads for a conjugate gradient solver, one thread per GPU
*
*/

#ifndef VIENNACL_WITH_OPENCL
  #define VIENNACL_WITH_OPENCL
#endif

// include necessary system headers
#include <iostream>

//
// ublas includes
//
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/triangular.hpp>
#include <boost/numeric/ublas/matrix_sparse.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/ublas/operation.hpp>
#include <boost/numeric/ublas/operation_sparse.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/lu.hpp>

// Must be set if you want to use ViennaCL algorithms on ublas objects
#define VIENNACL_WITH_UBLAS 1

//include basic scalar and vector types of ViennaCL
#include "viennacl/scalar.hpp"
#include "viennacl/vector.hpp"
#include "viennacl/compressed_matrix.hpp"
#include "viennacl/io/matrix_market.hpp"

#include "viennacl/ocl/device.hpp"
#include "viennacl/ocl/platform.hpp"
#include "viennacl/ocl/backend.hpp"

//include the generic inner product functions of ViennaCL
#include "viennacl/linalg/norm_2.hpp"
#include "viennacl/linalg/cg.hpp"

// Some helper functions for this tutorial:
#include "Random.hpp"
#include "vector-io.hpp"

using namespace boost::numeric;

#include <boost/thread.hpp>

template<typename NumericT>
class worker
{
public:
  worker(std::size_t tid) : thread_id_(tid) {}

  void operator()()
  {
    //
    // Set up some ublas objects
    //
    ublas::vector<NumericT> rhs;
    ublas::vector<NumericT> ref_result;
    ublas::compressed_matrix<NumericT> ublas_matrix;

    //
    // Read system from file
    //
    if (!viennacl::io::read_matrix_market_file(ublas_matrix, "../examples/testdata/mat65k.mtx"))
    {
      std::cout << "Error reading Matrix file" << std::endl;
      return;
    }

    if (!readVectorFromFile("../examples/testdata/rhs65025.txt", rhs))
    {
      std::cout << "Error reading RHS file" << std::endl;
      return;
    }

    if (!readVectorFromFile("../examples/testdata/result65025.txt", ref_result))
    {
      std::cout << "Error reading Result file" << std::endl;
      return;
    }

    //
    // Set up some ViennaCL objects in the respective context
    //
    viennacl::context ctx(viennacl::ocl::get_context(static_cast<long>(thread_id_)));

    std::size_t vcl_size = rhs.size();
    viennacl::compressed_matrix<NumericT> vcl_compressed_matrix(ctx);
    viennacl::vector<NumericT> vcl_rhs(vcl_size, ctx);
    viennacl::vector<NumericT> vcl_ref_result(vcl_size, ctx);

    viennacl::copy(rhs.begin(), rhs.end(), vcl_rhs.begin());
    viennacl::copy(ref_result.begin(), ref_result.end(), vcl_ref_result.begin());

    //
    // Transfer ublas-matrix to GPU:
    //
    viennacl::copy(ublas_matrix, vcl_compressed_matrix);

    viennacl::vector<NumericT> vcl_result = viennacl::linalg::solve(vcl_compressed_matrix, vcl_rhs, viennacl::linalg::cg_tag());

    std::stringstream ss;
    ss << "Result of thread " << thread_id_ << " on device " << viennacl::ocl::get_context(static_cast<long>(thread_id_)).devices()[0].name() << ": " << vcl_result[0] << ", should: " << ref_result[0] << std::endl;
    message_ = ss.str();
  }

  std::string message() const { return message_; }

private:
  std::string message_;
  std::size_t thread_id_;
};

int main()
{
  //Change this type definition to double if your gpu supports that
  typedef float       ScalarType;

  if (viennacl::ocl::get_platforms().size() == 0)
  {
    std::cerr << "Error: No platform found!" << std::endl;
    return EXIT_FAILURE;
  }

  //
  // Part 1: Setup first device for first context, second device for second context:
  //
  viennacl::ocl::platform pf = viennacl::ocl::get_platforms()[0];
  std::vector<viennacl::ocl::device> const & devices = pf.devices();

  // Set first device to first context:
  viennacl::ocl::setup_context(0, devices[0]);

  // Set second device for second context (use the same device for the second context if only one device available):
  if (devices.size() > 1)
    viennacl::ocl::setup_context(1, devices[1]);
  else
    viennacl::ocl::setup_context(1, devices[0]);

  //
  // Part 2: Now let two threads operate on two GPUs in parallel
  //

  worker<ScalarType> work_functor0(0);
  worker<ScalarType> work_functor1(1);
  boost::thread worker_thread_0(boost::ref(work_functor0));
  boost::thread worker_thread_1(boost::ref(work_functor1));

  worker_thread_0.join();
  worker_thread_1.join();

  std::cout << work_functor0.message() << std::endl;
  std::cout << work_functor1.message() << std::endl;

  std::cout << "!!!! TUTORIAL COMPLETED SUCCESSFULLY !!!!" << std::endl;

  return EXIT_SUCCESS;
}


// clang-format off
// MUST BE at the beginning before any other <cmath> include (e.g. in armadillo's headers)
#define _USE_MATH_DEFINES // required for Visual Studio
#include <cmath>
// clang-format on

#include "libKriging/Bench.hpp"
#include "libKriging/OrdinaryKriging.hpp"

#include <armadillo>
#include <optim.hpp>
#include <tuple>

// #include "libKriging/covariance.h"

LIBKRIGING_EXPORT Bench::Bench(int _n) {
  n = _n;
}

////////////////// LogLik /////////////////////
//' @ref https://github.com/cran/DiceKriging/blob/master/R/logLikFun.R
//  model@covariance <- vect2covparam(model@covariance, param)
//  model@covariance@sd2 <- 1		# to get the correlation matrix
//
//  aux <- covMatrix(model@covariance, model@X)
//
//  R <- aux[[1]]
//  T <- chol(R)
//
//  x <- backsolve(t(T), model@y, upper.tri = FALSE)
//  M <- backsolve(t(T), model@F, upper.tri = FALSE)
//  z <- compute.z(x=x, M=M, beta=beta)
//  sigma2.hat <- compute.sigma2.hat(z)
//  logLik <- -0.5*(model@n * log(2*pi*sigma2.hat) + 2*sum(log(diag(T))) + model@n)

////////////////// LogLikGrad /////////////////////
//' @ref https://github.com/cran/DiceKriging/blob/master/R/logLikGrad.R
//  logLik.derivative <- matrix(0,nparam,1)
//  x <- backsolve(T,z)			# compute x := T^(-1)*z
//  Rinv <- chol2inv(T)			# compute inv(R) by inverting T
//
//  Rinv.upper <- Rinv[upper.tri(Rinv)]
//  xx <- x%*%t(x)
//  xx.upper <- xx[upper.tri(xx)]
//
//  for (k in 1:nparam) {
//    gradR.k <- CovMatrixDerivative(model@covariance, X=model@X, C0=R, k=k)
//    gradR.k.upper <- gradR.k[upper.tri(gradR.k)]
//
//    terme1 <- sum(xx.upper*gradR.k.upper)   / sigma2.hat
//    # quick computation of t(x)%*%gradR.k%*%x /  ...
//    terme2 <- - sum(Rinv.upper*gradR.k.upper)
//    # quick computation of trace(Rinv%*%gradR.k)
//    logLik.derivative[k] <- terme1 + terme2
//  }

LIBKRIGING_EXPORT
arma::mat Bench::SolveTri(const arma::mat& Xtri, const arma::vec& y) {
  arma::mat s;
  for (int i = 0; i < n; i++) {
    s = arma::solve(arma::trimatu(Xtri), y, arma::solve_opts::fast);
  }
  return s;
}

LIBKRIGING_EXPORT
arma::mat Bench::CholSym(const arma::mat& Rsym) {
  arma::mat s;
  for (int i = 0; i < n; i++) {
    s = arma::chol(Rsym);
  }
  return s;
}

LIBKRIGING_EXPORT
std::tuple<arma::mat, arma::mat> Bench::QR(const arma::mat& M) {
  arma::mat Q;
  arma::mat R;
  for (int i = 0; i < n; i++) {
    arma::qr_econ(Q, R, M);
  }
  return std::make_tuple(std::move(Q), std::move(R));
}

LIBKRIGING_EXPORT
arma::mat Bench::InvSymPD(const arma::mat& Rsympd) {
  arma::mat s;
  for (int i = 0; i < n; i++) {
    s = arma::inv_sympd(Rsympd);
  }
  return s;
}

LIBKRIGING_EXPORT
double Bench::LogLik(OrdinaryKriging& ok, const arma::vec& theta) {
  // arma::vec theta = 0.5*ones(ok->X().n_cols)
  double s = 0;
  for (int i = 0; i < n; i++) {
    s += ok.logLikelihood(theta);
  }
  return s / n;
}

LIBKRIGING_EXPORT
arma::vec Bench::LogLikGrad(OrdinaryKriging& ok, const arma::vec& theta) {
  // arma::vec theta = 0.5*ones(ok->X().n_cols)
  arma::vec s = arma::zeros(theta.n_elem);
  for (int i = 0; i < n; i++) {
    s += ok.logLikelihoodGrad(theta);
  }
  return s / n;
  }

double a = .5;
double b = 50;
inline double rosenbrock_fun(arma::vec X) noexcept {
  return (a-X(0))*(a-X(0))+b*(X(1)-X(0)*X(0))*(X(1)-X(0)*X(0));
};
inline arma::vec rosenbrock_grad(arma::vec X) noexcept {
  arma::vec g = arma::zeros(2);
  g(0) = -2*(a-X(0)) + 4*b*(X(0)*X(0)*X(0) - X(1)*X(0));
  g(1) = 2*b*(X(1) - X(0)*X(0));
  return g;
};

  double ofn_rosenbrock(const arma::vec& x,arma::vec* grad_out,void* ofn_data) {
    if (ofn_data) {
    Bench::OFNData* fd = reinterpret_cast<Bench::OFNData*>(ofn_data);
    if (fd->histx.n_cols != x.n_elem)
      fd->histx = arma::reshape(x,1,x.n_elem);
    else
      fd->histx.insert_rows(fd->histx.n_rows,trans(arma::mat(x)));
    }
    arma::cout << "x "<< x << arma::endl;
    if (grad_out) {
      arma::cout << "g";
      *grad_out = rosenbrock_grad(x);
    } else     arma::cout << "o";

    return rosenbrock_fun(x);
  }

  LIBKRIGING_EXPORT
    double Bench::Rosenbrock(arma::vec& x) {
      return rosenbrock_fun(x);
    }

  LIBKRIGING_EXPORT
    arma::vec Bench::RosenbrockGrad(arma::vec& x) {
      return rosenbrock_grad(x);
    }

  LIBKRIGING_EXPORT
    arma::mat Bench::OptimRosenbrock(arma::vec& x0) {

      optim::algo_settings_t algo_settings;
      algo_settings.vals_bound = true;
      algo_settings.lower_bounds = arma::zeros<arma::vec>(2);
      algo_settings.upper_bounds = arma::ones<arma::vec>(2);

      algo_settings.iter_max = 10;  // TODO change by default?
      algo_settings.err_tol = 1e-9;

      algo_settings.gd_method = 2;
      algo_settings.gd_settings.step_size=0.01;
      algo_settings.gd_settings.norm_term=1E-7;
      algo_settings.gd_settings.ada_rho=0.9;

      algo_settings.cg_method = 2;

      Bench::OFNData ofn_data; // FIXME AFTER
      arma::cout << "> bfgs 2 ";
      ofn_data.histx = arma::zeros(1,2);
      bool bfgs_ok = optim::cg(
        x0,
        ofn_rosenbrock,
        (void*)(&ofn_data),
        algo_settings);
      arma::cout << " <" << arma::endl;

      return ofn_data.histx;
    }


#include "utils.h"
#include <NTL/GF2X.h>
#include <openssl/evp.h>
#include <openssl/sha.h>
#include <vector>
#include <rank/rank_params.h>
#include <random>

static int is_element_in_vector(const NTL::GF2E &element, const NTL::vec_GF2E &vector) {
    for (int i = 0; i < vector.length(); i++) {
        if (element == vector[i]) {
            return 1;
        }
    }
    return 0;
}

static void gf2_from_gf2e_element(NTL::vec_GF2 &res, const NTL::GF2E &v) {

    if (NTL::IsZero(v)) {
        for (int i = 0; i < v.degree(); i++) {
            res.append(NTL::GF2(0));
            return;
        }
    }

    for (int i = 0; i < v.degree(); i++) {
        res.append(v._GF2E__rep[i]);
    }
}

NTL::GF2X utils::gf2x_from_gf2(const NTL::Vec<NTL::GF2> &in) {

    NTL::GF2X v;
    v.SetLength(in.length());

    for (int i = 0; i < in.length(); i++) {
        NTL::SetCoeff(v, i, in[i]);
    }

    return v;
}

NTL::Vec<NTL::GF2> utils::gf2_from_gf2x(const NTL::GF2X &in, int len) {
    NTL::Vec<NTL::GF2> v;
    NTL::VectorCopy(v, in, len);

    return v;
}

NTL::GF2X utils::gf2x_from_matgf2(const NTL::mat_GF2 &in) {
    NTL::GF2X v;
    v.SetLength(in.NumRows() * in.NumCols());

    for (int i = 0; i < in.NumRows(); i++) {
        for (int j = 0; j < in.NumCols(); j++) {
            v[(i * in.NumCols()) + j] = in[i][j];
        }
    }
    return v;
}

NTL::GF2X utils::gf2x_from_gf2e(const NTL::vec_GF2E &in) {
    NTL::GF2X v;
    v.SetLength(in.length() * in[0].degree());

    for (int i = 0; i < in.length(); i++) {
        for (int j = 0; j < in[i].degree(); j++) {
            NTL::SetCoeff(v, (i * in[i].degree()) + j, in[i]._GF2E__rep[j]);
        }
    }

    return v;
}

NTL::vec_GF2 utils::gf2_from_gf2e(const NTL::vec_GF2E &in) {
    NTL::vec_GF2 v;
    v.SetLength(in.length() * in[0].degree());

    for (int i = 0; i < in.length(); i++) {
        for (int j = 0; j < in[i].degree(); j++) {
            v[(i * in[i].degree()) + j] = in[i]._GF2E__rep[j];
        }
    }

    return v;
}

NTL::vec_GF2E utils::gf2e_from_gf2x(const NTL::GF2X &in, int length) {
    NTL::vec_GF2E v;
    auto v_size = std::ceil((float) length / (float) v[0].degree());
    v.SetLength(v_size);

    for (int i = 0; i < v.length(); i++) {
        for (int j = 0; j < v[i].degree(); j++) {
            NTL::SetCoeff(v[i]._GF2E__rep, j, in[(i * v[i].degree()) + j]);
        }
    }
    return v;
}

NTL::vec_GF2 utils::encode(const NTL::mat_GF2 &P, const NTL::mat_GF2 &Q, const NTL::vec_GF2E &V,
                           int size) {

    uint digest_length = SHA512_DIGEST_LENGTH;

    const EVP_MD* algorithm = EVP_sha3_512();
    EVP_MD_CTX* context = EVP_MD_CTX_new();
    EVP_DigestInit_ex(context, algorithm, nullptr);

    NTL::vec_GF2 out;
    std::vector<uint8_t> hash(digest_length);

    auto P_gf2x = gf2x_from_matgf2(P);
    std::vector<uint8_t> P_bytes(NTL::NumBytes(P_gf2x));
    NTL::BytesFromGF2X(P_bytes.data(), P_gf2x, P_bytes.size());
    EVP_DigestUpdate(context, P_bytes.data(), P_bytes.size());

    auto Q_gf2x = gf2x_from_matgf2(Q);
    std::vector<uint8_t> Q_bytes(NTL::NumBytes(Q_gf2x));
    NTL::BytesFromGF2X(Q_bytes.data(), Q_gf2x, Q_bytes.size());
    EVP_DigestUpdate(context, Q_bytes.data(), Q_bytes.size());

    auto V_gf2x = gf2x_from_gf2e(V);
    std::vector<uint8_t> V_bytes(NTL::NumBytes(V_gf2x));
    NTL::BytesFromGF2X(V_bytes.data(), V_gf2x, V_bytes.size());
    EVP_DigestUpdate(context, V_bytes.data(), V_bytes.size());

    EVP_DigestFinal_ex(context, hash.data(), &digest_length);

    auto hash_gf2x = NTL::GF2XFromBytes(hash.data(), hash.size());

    out = gf2_from_gf2x(hash_gf2x, size);

    EVP_MD_CTX_destroy(context);

    return out;
}

NTL::vec_GF2 utils::encode_2(const NTL::vec_GF2E &V, int size) {

    uint digest_length = SHA512_DIGEST_LENGTH;
    const EVP_MD* algorithm = EVP_sha3_512();
    EVP_MD_CTX* context = EVP_MD_CTX_new();
    EVP_DigestInit_ex(context, algorithm, nullptr);

    NTL::vec_GF2 out;
    std::vector<uint8_t> hash(digest_length);

    auto V_gf2x = gf2x_from_gf2e(V);
    std::vector<uint8_t> V_bytes(NTL::NumBytes(V_gf2x));
    NTL::BytesFromGF2X(V_bytes.data(), V_gf2x, V_bytes.size());
    EVP_DigestUpdate(context, V_bytes.data(), V_bytes.size());

    EVP_DigestFinal_ex(context, hash.data(), &digest_length);

    auto hash_gf2x = NTL::GF2XFromBytes(hash.data(), hash.size());

    out = gf2_from_gf2x(hash_gf2x, size);

    EVP_MD_CTX_destroy(context);

    return out;
}

NTL::vec_GF2E utils::gf2e_from_two_gf2(const NTL::vec_GF2 &first, const NTL::vec_GF2 &second) {
    NTL::vec_GF2E out;
    auto out_size = std::ceil((float) (first.length() + second.length()) / (float) out[0].degree());
    out.SetLength(out_size);

    NTL::vec_GF2 first_second = first;
    first_second.append(second);

    for (int i = 0; i < out.length(); i++) {
        for (int j = 0; j < out[i].degree(); j++) {
            NTL::SetCoeff(out[i]._GF2E__rep, j, first_second[(i * out[i].degree()) + j]);
        }
    }

    return out;
}

NTL::vec_GF2E utils::gf2e_from_vec_gf2(const NTL::vec_GF2 &in) {
    NTL::vec_GF2E out;
    auto out_size = std::ceil((float) in.length() / (float) out[0].degree());
    out.SetLength(out_size);

    for (int i = 0; i < out.length(); i++) {
        for (int j = 0; j < out[i].degree(); j++) {
            NTL::SetCoeff(out[i]._GF2E__rep, j, in[(i * out[i].degree()) + j]);
        }
    }

    return out;
}

NTL::Vec<NTL::GF2> utils::encode_binary_vector(const NTL::Vec<NTL::GF2> &in, int size) {

    uint digest_length = SHA512_DIGEST_LENGTH;
    const EVP_MD* algorithm = EVP_sha3_512();
    EVP_MD_CTX* context = EVP_MD_CTX_new();
    EVP_DigestInit_ex(context, algorithm, nullptr);

    auto in_to_gf2x = utils::gf2x_from_gf2(in);
    std::vector<uint8_t> in_to_bytes(NTL::NumBytes(in_to_gf2x));
    NTL::BytesFromGF2X(in_to_bytes.data(), in_to_gf2x, in_to_bytes.size());

    std::vector<uint8_t> hash(NTL::NumBytes(in_to_gf2x));

    EVP_DigestUpdate(context, in_to_bytes.data(), in_to_bytes.size());
    EVP_DigestFinal_ex(context, hash.data(), &digest_length);

    NTL::GF2X hash_to_gf2x;
    NTL::GF2XFromBytes(hash_to_gf2x, hash.data(), hash.size());

    auto hash_gf2 = utils::gf2_from_gf2x(hash_to_gf2x, size);

    EVP_MD_CTX_destroy(context);

    return hash_gf2;
}

void utils::sample_messages(
        NTL::vec_GF2 &m_1,
        NTL::vec_GF2 &m_2,
        NTL::vec_GF2 &m_3,
        int size) {

    NTL::vec_GF2 _m_1, _m_2, _m_3;
    {
        _m_1.SetLength(4);
        _m_2.SetLength(4);
        _m_3.SetLength(4);

        _m_1[2] = 1;
        _m_1[3] = 1;

        _m_2[1] = 1;
        _m_2[3] = 1;

        _m_3[3] = 1;
    }

    m_1.append(_m_1);
    m_2.append(_m_2);
    m_3.append(_m_3);

    for (int i = 0; i < size - 1; i++) {

        std::vector<uint> permutation;
        {
            permutation.resize(4);
        }

        create_random_permutation(
                permutation,
                4);

        NTL::mat_GF2 permutation_matrix;
        {
            permutation_matrix.SetDims(4, 4);
        }

        for (uint j = 0; j < permutation.size(); j++) {
            permutation_matrix[j][permutation[j]] = NTL::GF2(1);
        }

        m_1.append(permutation_matrix * _m_1);
        m_2.append(permutation_matrix * _m_2);
        m_3.append(permutation_matrix * _m_3);

    }
}

void utils::create_random_permutation(
        std::vector<uint> &permutation,
        int size) {

    permutation.resize(size);

    for (uint i = 0; i < permutation.size(); i++) {
        permutation[i] = i;
    }

    auto shuffled_vector = shuffle(
            permutation,
            permutation.size());

    permutation = shuffled_vector;

}

void utils::create_permutation_matrix(
        NTL::mat_GF2 &P,
        const std::vector<uint> &permutation) {

    for (uint i = 0; i < permutation.size(); i++) {
        P[i][permutation[i]] = NTL::GF2(1);
    }
}

void utils::create_permutation_matrix(
        NTL::mat_GF2 &P,
        int size) {

    P.kill();
    P.SetDims(size, size);

    std::vector<uint> permutation(size);

    create_random_permutation(
            permutation,
            size);

    for (uint i = 0; i < permutation.size(); i++) {
        P[i][permutation[i]] = NTL::GF2(1);
    }
}

NTL::mat_GF2 utils::mat_gf2_from_vec_gf2e(const NTL::vec_GF2E &v) {
    NTL::mat_GF2 M;
    M.SetDims(v.length(), v[0].degree());

    for (int j = 0; j < v.length(); j++) {
        for (int z = 0; z < v[0].degree(); z++) {
            if (NTL::IsZero(v[j]._GF2E__rep[z])) {
                M[j][z] = NTL::GF2(0);
            } else {
                M[j][z] = NTL::GF2(1);
            }
        }
    }
    return M;
}

void utils::generate_relation_matrix(
        NTL::mat_GF2 &_R,
        NTL::Vec<NTL::mat_GF2> &R,
        const NTL::Vec<NTL::vec_GF2> &m_tilde,
        const NTL::Vec<NTL::vec_GF2> &m,
        int size) {

    _R.SetDims(size, 4 * size);

    NTL::vec_GF2 tmp;
    tmp.SetLength(_R.NumCols());

    for (int i = 0; i < _R.NumRows(); i++) {
        for (int j = 0; j < _R.NumCols(); j++) {
            _R[i][j] = NTL::GF2(1);
        }
    }

    for (int i = 0; i < m[0].length(); i++) {
        for (int j = 0; j < m.length() - 1; j++) {
            for (int q = 0; q < m_tilde[j].length(); q++) {
                if (NTL::IsZero(m[j][i]) && (NTL::IsZero(m_tilde[j][q]))) {
                    tmp[q] = NTL::GF2(1);
                }
                if (NTL::IsZero(m[j][i]) && (!NTL::IsZero(m_tilde[j][q]))) {
                    tmp[q] = NTL::GF2(0);
                }
                if (!NTL::IsZero(m[j][i]) && (NTL::IsZero(m_tilde[j][q]))) {
                    tmp[q] = NTL::GF2(0);
                }
                if (!NTL::IsZero(m[j][i]) && (!NTL::IsZero(m_tilde[j][q]))) {
                    tmp[q] = NTL::GF2(1);
                }
            }
            for (int c = 0; c < tmp.length(); c++) {
                _R[i][c] *= tmp[c];
            }
        }
    }

    NTL::mat_GF2 R_tmp;
    R_tmp.SetDims(_R.NumRows(), _R.NumCols());

    std::vector<uint> v;
    v.resize(size);

    std::vector<std::vector<uint>> indexes;
    indexes.resize(size);

    for (int i = 0; i < _R.NumRows(); i++) {
        for (int j = 0; j < _R.NumCols(); j++) {
            if (NTL::IsOne(_R[i][j])) {
                indexes[i].push_back(j);
            }
        }
    }

    auto rnd_element = indexes[0][NTL::RandomBnd(indexes[0].size() - 1)];
    v[0] = rnd_element;
    R_tmp[0][rnd_element] = NTL::GF2(1);

    for (int i = 1; i < _R.NumRows(); i++) {
        do {
            rnd_element = indexes[i][NTL::RandomBnd(indexes[i].size() - 1)];
        } while (std::find(v.begin(), v.end(), rnd_element) != v.end());
        v[i] = rnd_element;
        R_tmp[i][rnd_element] = NTL::GF2(1);
    }

    _R = R_tmp;

    // Generate R
    R.SetLength(J);

    for (int i = 0; i < J; i++) {
        R[i].SetDims(size, size);
        for (int j = 0; j < size; j++) {
            for (int z = 0; z < size; z++) {
                R[i][j][z] = _R[j][z + (i * size)];
            }
        }
    }
}

int utils::gauss_row_reduced_echelon_form(NTL::mat_GF2 &M) {
    long row;
    long max_i = (int) floor((M.NumRows() + (M.NumCols() - 1)) / M.NumCols());
    for (long i = 0; i < max_i; i++) {
        for (long j = 0; j < M.NumCols(); j++) {

            row = i * M.NumCols() + j;

            if (row >= M.NumRows()) {
                return 1;
            }
            for (long k = row + 1; k < M.NumRows(); k++) {
                if (!NTL::IsZero(M[row][j] + M[k][j])) {
                    M[row] += M[k];
                }
            }

            for (long k = 0; k < M.NumRows(); k++) {
                if (k != row) {
                    if (!NTL::IsZero(M[k][j])) {
                        M[k] += M[row];
                    }
                }
            }
        }
    }
    return 0;
}

int utils::gauss_row_reduced_echelon_form_gf2e(
        NTL::mat_GF2E &M) {
    int i = 0;
    int j = 0;
    int found = 0;

    for (;;) {
        if (!NTL::IsZero(M[i][j])) {
            auto element = M[i][j];
            for (int c = 0; c < M.NumCols(); ++c) {
                M[i][c] = M[i][c] / element;
            }
        }

        if (!NTL::IsZero(M[i][j])) {
            for (int k = 0; k < M.NumRows(); k++) {
                if (k != i) {
                    auto t = M[k][j] / M[i][j];
                    for (int h = 0; h < M.NumCols(); h++) {
                        M[k][h] = M[k][h] - M[i][h] * t;
                    }
                }
            }
            i++;
            j++;
        } else {
            found = 0;
            for (int k = i; k < M.NumRows(); k++) {
                if (!NTL::IsZero(M[k][j])) {
                    for (int h = 0; h < M.NumCols(); h++) {
                        M[i][h] = M[i][h] + M[k][h];
                    }
                    found = 1;
                    break;
                }
            }
            if (found == 0) {
                j++;
            }
        }
        if (i > (M.NumRows() - 1) || j > (M.NumCols() - 1)) {
            break;
        }
    }

    return 0;
}

int utils::convert_matrix_to_rref_and_check_solutions(NTL::mat_GF2 &M) {

    utils::gauss_row_reduced_echelon_form(M);

    for (int i = 0; i < M.NumRows(); i++) {
        if (NTL::weight(NTL::IsZero(M[i]))) {
            break;
        }
        if (NTL::IsOne(M[i][M.NumCols() - 1]) && (NTL::weight(M[i]) == 1)) {
            return 1;
        }
    }
    return 0;
}

int utils::convert_matrix_to_rref_and_check_solutions(NTL::mat_GF2E &M) {

    utils::gauss_row_reduced_echelon_form_gf2e(M);

    for (int i = 0; i < K; i++) {
        for (int j = 0; j < K; j++) {
            if (i == j) {
                if (M[i][j]._GF2E__rep != NTL::GF2(1)) {
                    return 1;
                }
            } else {
                if (M[i][j]._GF2E__rep != NTL::GF2(0)) {
                    return 1;
                }
            }
        }
    }

    return 0;
}

void utils::generate_random_matrix_gf2e(
        NTL::mat_GF2E &g,
        int num_of_rows,
        int num_of_cols) {

    g.kill();
    g = NTL::random_mat_GF2E(num_of_rows, num_of_cols);
}

void utils::generate_random_square_invertible_matrix(
        NTL::mat_GF2 &M,
        int size) {

    do {
        M = NTL::random_mat_GF2(size, size);
    } while (NTL::determinant(M) == 0);
}

void utils::generate_random_binary_matrix(
        NTL::mat_GF2 &A,
        int num_of_rows,
        int num_of_cols) {

    A.kill();
    A = NTL::random_mat_GF2(num_of_rows, num_of_cols);
}

void utils::generate_random_binary_vector_gf2e(
        NTL::vec_GF2E &v,
        int length) {

    v.SetLength(length);
    v = NTL::random_vec_GF2E(length);
}

void utils::generate_vector_of_specific_rank(
        NTL::vec_GF2E &e,
        int length,
        int rank) {

    e.SetLength(length);
    auto degree = e[0].degree();

    NTL::mat_GF2 _e;
    _e.SetDims(degree, length);

    NTL::vec_GF2E F;
    F.append(NTL::GF2E());

    int dim = 0;

    while (dim < rank) {
        auto n = NTL::random_GF2E();
        if (!is_element_in_vector(n, F)) {
            NTL::GF2E tmp;
            auto size_of_F = F.length();
            for (int i = 0; i < size_of_F; i++) {
                tmp = n + F[i];
                F.append(tmp);
            }
            dim++;
        }
    }

    NTL::mat_GF2 tmp_e;
    uint _rank = 0;
    do {
        for (int i = 0; i < length; i++) {
            auto rnd_element = F[NTL::RandomBnd(F.length() - 1)];
            NTL::vec_GF2 v;
            gf2_from_gf2e_element(v, rnd_element);

            for (int c = 0; c < v.length(); c++) {
                _e[c][i] = v[c];
            }
        }

        for (int i = 0; i < length; i++) {
            for (int j = 0; j < degree; j++) {
                NTL::SetCoeff(e[i].LoopHole(), j, _e[j][i]);
            }
        }

        tmp_e = utils::mat_gf2_from_vec_gf2e(e);
        _rank = NTL::gauss(tmp_e);
    } while (_rank != rank);

}

void utils::generate_random_binary_vector(
        NTL::vec_GF2 &v,
        int length) {

    v.SetLength(length);
    v = NTL::random_vec_GF2(length);
}

void utils::generate_vector_of_weight_w(
        NTL::vec_GF2 &e,
        int length,
        int weight) {

    e.SetLength(length);

    std::vector<uint> pos;
    for (long i = 0; i < e.length(); i++) {
        pos.push_back(i);
    }

    auto shuffled_vector = shuffle(
            pos,
            pos.size());

    for (long i = 0; i < weight; i++) {
        e[shuffled_vector[i]] = NTL::GF2(1);
    }
}

int utils::solve_equation(
        NTL::vec_GF2 &res,
        const NTL::mat_GF2 &M,
        const NTL::vec_GF2 &v) {

    auto A_T = NTL::transpose(M);

    A_T._mat__rep.append(v);
    A_T = NTL::transpose(A_T);

    if (utils::convert_matrix_to_rref_and_check_solutions(A_T) != 0) {
        return 1;
    }

    A_T = NTL::transpose(A_T);
    res = A_T[A_T.NumRows() - 1];

    return 0;
}

int utils::solve_equation(
        NTL::vec_GF2E &res,
        const NTL::mat_GF2E &M,
        const NTL::vec_GF2E &v) {

    auto A_T = M;

    A_T._mat__rep.append(v);
    A_T = NTL::transpose(A_T);

    if (utils::convert_matrix_to_rref_and_check_solutions(A_T) != 0) {
        return 1;
    }

    A_T = NTL::transpose(A_T);
    res = A_T[A_T.NumRows() - 1];

    return 0;
}

int utils::rank_of_vector(
        NTL::vec_GF2E &v) {

    NTL::mat_GF2 M;
    M.SetDims(v.length(), v[0].degree());

    for (int z = 0; z < EN; z++) {
        for (int j = 0; j < EM; j++) {
            if (NTL::IsZero(v[z]._GF2E__rep[j])) {
                M[z][j] = NTL::GF2(0);
            } else {
                M[z][j] = NTL::GF2(1);
            }
        }
    }

    auto _M = NTL::transpose(M);
    return NTL::gauss(_M);
}

std::vector<uint> utils::shuffle(
        const std::vector<uint> &input,
        int array_size) {

    std::vector<uint> index_array(input.size());
    std::vector<uint> output(input.size());

    int index;
    for (int i = 0; i < array_size; i++) {
        do {
            index = rand() % array_size;
        } while (index_array[index] != 0);
        index_array[index] = 1;
        output[i] = input[index];
    }

    return output;
}


#pragma once

#include <iostream>
#include <unistd.h>

#include <Eigen/Core>
#include <Eigen/Geometry>
#include <Eigen/Cholesky>
#include <Eigen/StdVector>

#include <opencv2/opencv.hpp>

/**
 * @file   types.h
 * @Author Giorgio Grisetti
 * @date   December 2014
 * @brief  This file contains some useful defines about
 * Eigen types, common Mat opencv specializations, and transforms mappings
 */

namespace srrg_core {

  //!a vector of Vector2f with alignment
  typedef std::vector<Eigen::Vector2i, Eigen::aligned_allocator<Eigen::Vector2i> > Vector2iVector;

  //!a vector of Vector3f with alignment
  typedef std::vector<Eigen::Vector3f, Eigen::aligned_allocator<Eigen::Vector3f> > Vector3fVector;

  //!a vector of Vector2f with alignment
  typedef std::vector<Eigen::Vector2f, Eigen::aligned_allocator<Eigen::Vector2f> > Vector2fVector;

  //!a vector of Matrix3f with alignment
  typedef std::vector<Eigen::Matrix3f, Eigen::aligned_allocator<Eigen::Matrix3f> > Matrix3fVector;

  //!a vector of Matrix2f with alignment
  typedef std::vector<Eigen::Matrix2f, Eigen::aligned_allocator<Eigen::Matrix2f> > Matrix2fVector;

  //!a 4x3 float matrix
  typedef Eigen::Matrix<float, 4, 3> Matrix4_3f;

  //!a 4x6 float matrix
  typedef Eigen::Matrix<float, 4, 6> Matrix4_6f;

  //!a 2x6 float matrix
  typedef Eigen::Matrix<float, 2, 6> Matrix2_6f;

  //!a 3x6 float matrix
  typedef Eigen::Matrix<float, 3, 6> Matrix3_6f;

  //!a 5x5 float matrix
  typedef Eigen::Matrix<float, 5, 5> Matrix5f;

  //!a 6x6 float matrix
  typedef Eigen::Matrix<float, 6, 6> Matrix6f;

  //!a 6   float vector
  typedef Eigen::Matrix<float, 6, 1> Vector6f;

  //!a 9x6 float matrix
  typedef Eigen::Matrix<float, 9, 6> Matrix9_6f;

  //!a 9x9 float matrix
  typedef Eigen::Matrix<float, 9, 9> Matrix9f;

  //!a 9 float vector
  typedef Eigen::Matrix<float, 9, 1> Vector9f;

  //!a 6x3 float matrix
  typedef Eigen::Matrix<float, 6, 3> Matrix6_3f;

  //!a 5 float cv vector
  typedef cv::Vec<float, 5> Vec5f;

  //!a 7 float cv vector
  typedef cv::Vec<float, 7> Vec7f;

  //!check if an Eigen type contains a nan element
  //!@returns true if at least one element of
  //!the argument is null
  template <class T>
  bool isNan(const T& m){
    for (int i=0; i< m.rows(); i++) {
      for (int j=0; j< m.cols(); j++) {
	float v = m(i,j);
	if ( std::isnan( v ) )
	  return true;
      }
    }
    return false;
  }

  //!converts from 6 vector to isometry
  //!@param t: a vector (tx, ty, tz, qx, qy, qz) reptesenting the transform.
  //!(qx, qy, qz) are the imaginary part of a normalized queternion, with qw>0.
  //!@returns the isometry corresponding to the transform described by t
  inline Eigen::Isometry3f v2t(const Vector6f& t){
    Eigen::Isometry3f T;
    T.setIdentity();
    T.translation()=t.head<3>();
    float w=t.block<3,1>(3,0).squaredNorm();
    if (w<1) {
      w=sqrt(1-w);
      T.linear()=Eigen::Quaternionf(w, t(3), t(4), t(5)).toRotationMatrix();
    } else {
      Eigen::Vector3f q=t.block<3,1>(3,0);
      q.normalize();
      T.linear()=Eigen::Quaternionf(0, q(0), q(1), q(2)).toRotationMatrix();
    }
    return T;
  }
  inline Eigen::Isometry3d v2t(const Eigen::Matrix<double, 6, 1>& t){
    Eigen::Isometry3d T;
    T.setIdentity();
    T.translation()=t.head<3>();
    float w=t.block<3,1>(3,0).squaredNorm();
    if (w<1) {
      w=sqrt(1-w);
      T.linear()=Eigen::Quaterniond(w, t(3), t(4), t(5)).toRotationMatrix();
    } else {
      Eigen::Vector3d q=t.block<3,1>(3,0);
      q.normalize();
      T.linear()=Eigen::Quaterniond(0, q(0), q(1), q(2)).toRotationMatrix();
    }
    return T;
  }

  //!converts from isometry to 6 vector
  //!@param t: an isometry
  //!@returns a vector (tx, ty, tz, qx, qy, qz) reptesenting the transform.
  //!(qx, qy, qz) are the imaginary part of a normalized queternion, with qw>0.
  inline Vector6f t2v(const Eigen::Isometry3f& t){
    Vector6f v;
    v.head<3>()=t.translation();
    Eigen::Quaternionf q(t.linear());
    v(3) = q.x();
    v(4) = q.y();
    v(5) = q.z();
    if (q.w()<0)
      v.block<3,1>(3,0) *= -1.0f;
    return v;
  }

  //!converts from isometry2f to isometry3f
  //!@param t: an isometry2f
  //!@returns an isometry3f
  inline Eigen::Isometry3f toIsometry3f(const Eigen::Isometry2f& isometry2f){
    Eigen::Isometry3f isometry3f;
    isometry3f.linear().block<2,2>(0,0) = isometry2f.linear();
    isometry3f.translation().head<2>() = isometry2f.translation();
    return isometry3f;
  }

  //!computes the cross product matrix of the vector argument
  //!@param p: the vector
  //!@returns a 3x3 matrix
  inline Eigen::Matrix3f skew(const Eigen::Vector3f& p){
    Eigen::Matrix3f s;
    s <<
      0,  -p.z(), p.y(),
      p.z(), 0,  -p.x(),
      -p.y(), p.x(), 0;
    return s;
  }
  inline Eigen::Matrix3d skew(const Eigen::Vector3d& p){
    Eigen::Matrix3d s;
    s <<
      0,  -p.z(), p.y(),
      p.z(), 0,  -p.x(),
      -p.y(), p.x(), 0;
    return s;
  }

  inline Eigen::Isometry2f v2t(const Eigen::Vector3f& t){
    Eigen::Isometry2f T;
    T.setIdentity();
    T.translation()=t.head<2>();
    float c = cos(t(2));
    float s = sin(t(2));
    T.linear() << c, -s, s, c;
    return T;
  }

  inline Eigen::Vector3f t2v(const Eigen::Isometry2f& t){
    Eigen::Vector3f v;
    v.head<2>()=t.translation();
    v(2) = atan2(t.linear()(1,0), t.linear()(0,0));
    return v;
  }

  inline Eigen::Matrix3f Rx(float rot_x){
    float c=cos(rot_x);
    float s=sin(rot_x);
    Eigen::Matrix3f R;
    R << 1,  0, 0,
      0,  c,  -s,
      0,  s,  c;
    return R;
  }
  inline Eigen::Matrix3d Rx(const double& rot_x){
    const double c=cos(rot_x);
    const double s=sin(rot_x);
    Eigen::Matrix3d R;
    R << 1,  0, 0,
      0,  c,  -s,
      0,  s,  c;
    return R;
  }

  inline Eigen::Matrix3f Ry(float rot_y){
    float c=cos(rot_y);
    float s=sin(rot_y);
    Eigen::Matrix3f R;
    R << c,  0,  s,
      0 , 1,  0,
      -s,  0, c;
    return R;
  }
  inline Eigen::Matrix3d Ry(const double& rot_y){
    const double c=cos(rot_y);
    const double s=sin(rot_y);
    Eigen::Matrix3d R;
    R << c,  0,  s,
      0 , 1,  0,
      -s,  0, c;
    return R;
  }

  inline Eigen::Matrix3f Rz(float rot_z){
    float c=cos(rot_z);
    float s=sin(rot_z);
    Eigen::Matrix3f R;
    R << c,  -s,  0,
      s,  c,  0,
      0,  0,  1;
    return R;
  }
  inline Eigen::Matrix3d Rz(const double& rot_z){
    const double c=cos(rot_z);
    const double s=sin(rot_z);
    Eigen::Matrix3d R;
    R << c,  -s,  0,
      s,  c,  0,
      0,  0,  1;
    return R;
  }

  inline Eigen::Isometry3f v2tEuler(const Vector6f& v){
    Eigen::Isometry3f T;
    T.linear()=Rx(v[3])*Ry(v[4])*Rz(v[5]);
    T.translation()=v.head<3>();
    return T;
  }
  inline Eigen::Isometry3d v2tEuler(const Eigen::Matrix<double, 6, 1>& v){
    Eigen::Isometry3d T;
    T.linear()=Rx(v[3])*Ry(v[4])*Rz(v[5]);
    T.translation()=v.head<3>();
    return T;
  }

  inline Eigen::Matrix2f skew(const Eigen::Vector2f& p){
    Eigen::Matrix2f s;
    s <<
      0,  -p.y(),
      p.x(), 0;
    return s;
  }

  /** \typedef UnsignedCharImage
   * \brief An unsigned char cv::Mat.
   */
  typedef cv::Mat_<unsigned char> UnsignedCharImage;

  /** \typedef CharImage
   * \brief A char cv::Mat.
   */
  typedef cv::Mat_<char> CharImage;

  /** \typedef UnsignedShortImage
   * \brief An unsigned short cv::Mat.
   */
  typedef cv::Mat_<unsigned short> UnsignedShortImage;

  /** \typedef UnsignedIntImage
   * \brief An unsigned int cv::Mat.
   */
  typedef cv::Mat_<unsigned int> UnsignedIntImage;

  /** \typedef IntImage
   * \brief An int cv::Mat.
   */
  typedef cv::Mat_<int> IntImage;

  /** \typedef Int4Image
   * \brief A 4D int cv::Mat.
   */
  typedef cv::Mat_<cv::Vec4i> Int4Image;

  /** \typedef IntervalImage
   * \brief A 4D int cv::Mat used to store intervals.
   */
  typedef Int4Image IntervalImage;

  /** \typedef FloatImage
   * \brief A float cv::Mat.
   */
  typedef cv::Mat_<float> FloatImage;

  /** \typedef Float3Image
   * \brief A 3D float cv::Mat.
   */
  typedef cv::Mat_<cv::Vec3f> Float3Image;

  /** \typedef Float5Image
   * \brief A 5D float cv::Mat.
   */
  typedef cv::Mat_<Vec5f> Float5Image;

  /** \typedef Float5Image
   * \brief A 7D float cv::Mat.
   */
  typedef cv::Mat_<Vec7f> Float7Image;

  /** \typedef DoubleImage
   * \brief A double cv::Mat.
   */
  typedef cv::Mat_<double> DoubleImage;

  /** \typedef RawDepthImage
   * \brief An unsigned char cv::Mat used to for depth images with depth values expressed in millimeters.
   */
  typedef UnsignedShortImage RawDepthImage;

  /** \typedef IndexImage
   * \brief An int cv::Mat used to save the indeces of the points of a depth image inside a vector of points.
   */
  typedef IntImage IndexImage;

  /** \typedef DepthImage
   * \brief A float cv::Mat used to for depth images with depth values expressed in meters.
   */
  typedef cv::Mat_<cv::Vec3b> RGBImage;

  /** used to represent rgb values
   */
  typedef std::vector<cv::Vec3b> RGBVector;

  typedef std::vector<int> IntVector;

  typedef std::vector<float> FloatVector;

  typedef std::vector<std::pair<int, int> > IntPairVector;

  //ds overloaded opencv/eigen converters: double
  inline cv::Mat_<double> toCv(const Eigen::Matrix<double, 3, 3>& matrix_eigen_) {
    cv::Mat_<double> matrix_opencv(3, 3);
    for(uint32_t u = 0; u < 3; ++u) {
      for(uint32_t v = 0; v < 3; ++v) {
        matrix_opencv.at<double>(u, v) = matrix_eigen_(u, v);
      }
    }
    return matrix_opencv;
  }
  inline cv::Mat_<double> toCv(const Eigen::Matrix<double, 3, 4>& matrix_eigen_) {
    cv::Mat_<double> matrix_opencv(3, 4);
    for(uint32_t u = 0; u < 3; ++u) {
      for(uint32_t v = 0; v < 4; ++v) {
        matrix_opencv.at<double>(u, v) = matrix_eigen_(u, v);
      }
    }
    return matrix_opencv;
  }
  inline cv::Mat_<double> toCv(const Eigen::Matrix<double, 5, 1>& vector_eigen_) {
    cv::Mat_<double> vector_opencv(5,1);
    for(uint32_t u = 0; u < 5; ++u) {
      vector_opencv.at<double>(u) = vector_eigen_(u);
    }
    return vector_opencv;
  }
  inline Eigen::Matrix<double, 3, 1> fromCv(const cv::Vec<double, 3>& vector_opencv_) {
    Eigen::Matrix<double, 3, 1> vector_eigen;
    for(uint32_t u = 0; u < 3; ++u) {
      vector_eigen(u) = vector_opencv_(u);
    }
    return vector_eigen;
  }

  //ds overloaded opencv/eigen converters: float
  inline cv::Mat_<float> toCv(const Eigen::Matrix<float, 3, 3>& matrix_eigen_) {
    cv::Mat_<float> matrix_opencv(3, 3);
    for(uint32_t u = 0; u < 3; ++u) {
      for(uint32_t v = 0; v < 3; ++v) {
        matrix_opencv.at<float>(u, v) = matrix_eigen_(u, v);
      }
    }
    return matrix_opencv;
  }
  inline cv::Mat_<float> toCv(const Eigen::Matrix<float, 3, 4>& matrix_eigen_) {
    cv::Mat_<float> matrix_opencv(3, 4);
    for(uint32_t u = 0; u < 3; ++u) {
      for(uint32_t v = 0; v < 4; ++v) {
        matrix_opencv.at<float>(u, v) = matrix_eigen_(u, v);
      }
    }
    return matrix_opencv;
  }
  inline cv::Mat_<float> toCv(const Eigen::Matrix<float, 5, 1>& vector_eigen_) {
    cv::Mat_<float> vector_opencv(5,1);
    for(uint32_t u = 0; u < 5; ++u) {
      vector_opencv.at<float>(u) = vector_eigen_(u);
    }
    return vector_opencv;
  }
  inline Eigen::Matrix<float, 3, 1> fromCv(const cv::Vec<float, 3>& vector_opencv_) {
    Eigen::Matrix<float, 3, 1> vector_eigen;
    for(uint32_t u = 0; u < 3; ++u) {
      vector_eigen(u) = vector_opencv_(u);
    }
    return vector_eigen;
  }

  //ml std vector to Eigen converters
  inline Eigen::Matrix<float, 3, 1> fromFloatVector3f(const FloatVector& float_vector_){
    Eigen::Matrix<float, 3, 1> vector_eigen;
    assert(float_vector_.size() == 3);
    for (size_t i=0; i<3; i++)
      vector_eigen[i] = float_vector_[i];

    return vector_eigen;
  }
  inline FloatVector toFloatVector3f(const Eigen::Matrix<float, 3, 1>& vector_eigen_){
    FloatVector float_vector;
    float_vector.resize(3);
    for (size_t i=0; i<3; i++)
      float_vector[i] = vector_eigen_[i];
    return float_vector;
  }
}


/* -*- mode: c++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4 -*- */

/*
 Copyright (C) 2013, 2015 Peter Caspers

 This file is part of QuantLib, a free-software/open-source library
 for financial quantitative analysts and developers - http://quantlib.org/

 QuantLib is free software: you can redistribute it and/or modify it
 under the terms of the QuantLib license.  You should have received a
 copy of the license along with this program; if not, please email
 <quantlib-dev@lists.sf.net>. The license is also available online at
 <http://quantlib.org/license.shtml>.

 This program is distributed in the hope that it will be useful, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 FOR A PARTICULAR PURPOSE.  See the license for more details.
*/

/*! \file gsr.hpp
    \brief GSR 1 factor model
*/

#ifndef quantlib_gsr_hpp
#define quantlib_gsr_hpp

#include <ql/time/schedule.hpp>
#include <ql/math/integrals/simpsonintegral.hpp>
#include <ql/math/integrals/gausslobattointegral.hpp>
#include <ql/math/distributions/normaldistribution.hpp>

#include <ql/models/shortrate/onefactormodels/gaussian1dmodel.hpp>
#include <ql/processes/gsrprocess.hpp>

#include <boost/math/special_functions.hpp>

namespace QuantLib {

//! One factor gsr model, formulation is in forward measure

class Gsr : public Gaussian1dModel, public CalibratedModel {

  public:
    // constant mean reversion
    Gsr(const Handle<YieldTermStructure> &termStructure,
        const std::vector<Date> &volstepdates,
        const std::vector<Real> &volatilities, const Real reversion,
        const Real T = 60.0);
    // piecewise mean reversion (with same step dates as volatilities)
    Gsr(const Handle<YieldTermStructure> &termStructure,
        const std::vector<Date> &volstepdates,
        const std::vector<Real> &volatilities,
        const std::vector<Real> &reversions, const Real T = 60.0);
    // constant mean reversion with floating model data
    Gsr(const Handle<YieldTermStructure> &termStructure,
        const std::vector<Date> &volstepdates,
        const std::vector<Handle<Quote> > &volatilities,
        const Handle<Quote> reversion, const Real T = 60.0);
    // piecewise mean reversion with floating model data
    Gsr(const Handle<YieldTermStructure> &termStructure,
        const std::vector<Date> &volstepdates,
        const std::vector<Handle<Quote> > &volatilities,
        const std::vector<Handle<Quote> > &reversions, const Real T = 60.0);

    const Real numeraireTime() const;
    const void numeraireTime(const Real T);

    const Array &reversion() const { return reversion_.params(); }
    const Array &volatility() const { return sigma_.params(); }

    // calibration constraints

    Disposable<std::vector<bool> > FixedReversions() {
        std::vector<bool> res(reversions_.size(), true);
        std::vector<bool> vol(volatilities_.size(), false);
        res.insert(res.end(), vol.begin(), vol.end());
        return res;
    }

    Disposable<std::vector<bool> > FixedVolatilities() {
        std::vector<bool> res(reversions_.size(), false);
        std::vector<bool> vol(volatilities_.size(), true);
        res.insert(res.end(), vol.begin(), vol.end());
        return res;
    }

    Disposable<std::vector<bool> > MoveVolatility(Size i) {
        QL_REQUIRE(i < volatilities_.size(),
                   "volatility with index " << i << " does not exist (0..."
                                            << volatilities_.size() - 1 << ")");
        std::vector<bool> res(reversions_.size() + volatilities_.size(),
                              true);
        res[reversions_.size() + i] = false;
        return res;
    }

    Disposable<std::vector<bool> > MoveReversion(Size i) {
        QL_REQUIRE(i < reversions_.size(),
                   "reversion with index " << i << " does not exist (0..."
                                           << reversions_.size() - 1 << ")");
        std::vector<bool> res(reversions_.size() + volatilities_.size(),
                              true);
        res[i] = false;
        return res;
    }

    // With fixed reversion calibrate the volatilities one by one
    // to the given helpers. It is assumed that that volatility step
    // dates are suitable for this, i.e. they should be identical to
    // the fixing dates of the helpers (except for the last one where
    // we do not need a step). Also note that the endcritera reflect
    // only the status of the last calibration when using this method.
    void calibrateVolatilitiesIterative(
        const std::vector<boost::shared_ptr<CalibrationHelper> > &helpers,
        OptimizationMethod &method, const EndCriteria &endCriteria,
        const Constraint &constraint = Constraint(),
        const std::vector<Real> &weights = std::vector<Real>()) {

        for (Size i = 0; i < helpers.size(); i++) {
            std::vector<boost::shared_ptr<CalibrationHelper> > h(1, helpers[i]);
            calibrate(h, method, endCriteria, constraint, weights,
                      MoveVolatility(i));
        }
    }

    // With fixed volatility calibrate the reversions one by one
    // to the given helpers. In this case the step dates must be chosen
    // according to the maturities of the calibration instruments.
    void calibrateReversionsIterative(
        const std::vector<boost::shared_ptr<CalibrationHelper> > &helpers,
        OptimizationMethod &method, const EndCriteria &endCriteria,
        const Constraint &constraint = Constraint(),
        const std::vector<Real> &weights = std::vector<Real>()) {

        for (Size i = 0; i < helpers.size(); i++) {
            std::vector<boost::shared_ptr<CalibrationHelper> > h(1, helpers[i]);
            calibrate(h, method, endCriteria, constraint, weights,
                      MoveReversion(i));
        }
    }

  protected:
    const Real numeraireImpl(const Time t, const Real y,
                             const Handle<YieldTermStructure> &yts) const;

    const Real zerobondImpl(const Time T, const Time t, const Real y,
                            const Handle<YieldTermStructure> &yts) const;

    void generateArguments() {
        boost::static_pointer_cast<GsrProcess>(stateProcess_)->flushCache();
        notifyObservers();
    }

    void update() { LazyObject::update(); }

    void performCalculations() const {
        Gaussian1dModel::performCalculations();
        updateTimes();
        updateState();
    }

  private:
    void updateTimes() const;
    void updateState() const;
    void initialize(Real);

    Parameter &reversion_, &sigma_;

    std::vector<Handle<Quote> > volatilities_;
    std::vector<Handle<Quote> > reversions_;
    std::vector<Date> volstepdates_; // this is shared between vols and reversions
                                     // in case of piecewise reversions
    mutable std::vector<Time> volsteptimes_;
    mutable Array volsteptimesArray_; // FIXME this is redundant (just a copy of
                                      // volsteptimes_)
};

inline const Real Gsr::numeraireTime() const {
    return boost::dynamic_pointer_cast<GsrProcess>(stateProcess_)
        ->getForwardMeasureTime();
}

inline const void Gsr::numeraireTime(const Real T) {
    boost::dynamic_pointer_cast<GsrProcess>(stateProcess_)
        ->setForwardMeasureTime(T);
}
}

#endif