/*! \file HyParEstimator.cpp
    \brief Contains definitions of classes implementing hyper parameter estimation for DWI GPs

    \author Jesper Andersson
    \version 1.0b, Nov., 2013.
*/
// Definitions of classes implementing hyper parameter estimation for DWI GPs
//
// HyParEstimator.cpp
//
// Jesper Andersson, FMRIB Image Analysis Group
//
// Copyright (C) 2013 University of Oxford
//

#include <cstdlib>
#include <string>
#include <exception>
#include <vector>
#include <cmath>
#include <limits>
#include <ctime>
#include "armawrap/newmat.h"
#include "miscmaths/miscmaths.h"
#include "EddyHelperClasses.h"
#include "EddyUtils.h"
#include "HyParEstimator.h"

using namespace std;
using namespace EDDY;

DataSelector& DataSelector::ConcatToMe(const DataSelector& ds) EddyTry
{
  this->_data.reserve(this->_data.size()+ds._data.size());
  this->_coords.reserve(this->_coords.size()+ds._coords.size());
  this->_data.insert(this->_data.end(),ds._data.begin(),ds._data.begin());
  this->_coords.insert(this->_coords.end(),ds._coords.begin(),ds._coords.begin());
  return(*this);
} EddyCatch

NEWMAT::Matrix DataSelector::AsNEWMAT() const EddyTry
{
  NEWMAT::Matrix m(_data[0].Nrows(),_data.size());
  for (unsigned int i=0; i<_data.size(); i++) {
    m.Column(i+1) = _data[i];
  }
  return(m);
} EddyCatch

NEWMAT::Matrix DataSelector::CoordsAsNEWMAT() const EddyTry
{
  NEWMAT::Matrix m(4,_coords.size());
  for (unsigned int i=0; i<_coords.size(); i++) {
    m(1,i+1) = _coords[i][0]; m(2,i+1) = _coords[i][1]; m(3,i+1) = _coords[i][2]; m(4,i+1) = _coords[i][3];
  }
  return(m);
} EddyCatch

void DataSelector::common_constructor(const std::vector<std::shared_ptr<NEWIMAGE::volume<float> > >& idata,
				      const NEWIMAGE::volume<float>&                                 mask,
				      unsigned int                                                   rnvox,
				      float                                                          fwhm,
				      unsigned int                                                   d4,
				      int                                                            rndinit) EddyTry
{
  // Initialise rand if no explicit seed given
  if (rndinit) srand(rndinit);
  else srand(time(NULL));
  // Check that there are at least rnvox voxels available
  unsigned int nnz = 0; // Number of non-zero voxels in mask
  for (auto it = mask.fbegin(); it!=mask.fend(); ++it) nnz += (*it) ? 1 : 0;
  if (nnz < rnvox) throw EddyException("DataSelector::common_constructor: rnvox greater than number of non-zero voxels in mask.");
  // Select rnvox voxel coordinates within mask
  unsigned int nx = static_cast<unsigned int>(mask.xsize());
  unsigned int ny = static_cast<unsigned int>(mask.ysize());
  unsigned int nz = static_cast<unsigned int>(mask.zsize());
  std::vector<std::vector<unsigned int> > coords(rnvox);
  unsigned int nvox = 0;
  if (rnvox) {
    unsigned int l=0; unsigned int maxtry=1e8;
    for (l=0; l<maxtry; l++) {
      std::vector<unsigned int> nc(4);
      nc[0] = rand() % nx; nc[1] = rand() % ny; nc[2] = rand() % nz; nc[3] = d4;
      if (mask(nc[0],nc[1],nc[2])) {
	unsigned int vox;
	for (vox=0; vox<nvox; vox++) if (coords[vox]==nc) break;
	if (vox == nvox) coords[nvox++] = nc; // If no duplicate found
      }
      if (nvox == rnvox) break;
    }
    if (l==maxtry) throw EddyException("DataSelector::common_constructor: unable to find requested number of unique voxels");
  }
  std::vector<NEWMAT::ColumnVector> kernels(3);
  // Make convolution kernels if fwhm>0
  if (fwhm > 0.0) {
    float sx = (fwhm/std::sqrt(8.0*std::log(2.0)))/idata[0]->xdim();
    float sy = (fwhm/std::sqrt(8.0*std::log(2.0)))/idata[0]->ydim();
    float sz = (fwhm/std::sqrt(8.0*std::log(2.0)))/idata[0]->zdim();
    int nx=((int) (sx-0.001))*2 + 3;
    int ny=((int) (sy-0.001))*2 + 3;
    int nz=((int) (sz-0.001))*2 + 3;
    kernels[0] = NEWIMAGE::gaussian_kernel1D(sx,nx);
    kernels[1] = NEWIMAGE::gaussian_kernel1D(sy,ny);
    kernels[2] = NEWIMAGE::gaussian_kernel1D(sz,nz);
  }

  // Get time-series from the selected coordinates
  NEWMAT::ColumnVector vec(idata.size());
  _data.resize(nvox,vec);
  for (unsigned int i=0; i<idata.size(); i++) {
    const NEWIMAGE::volume<float>& vol = *idata[i];
    for (unsigned int j=0; j<nvox; j++) {
      if (fwhm > 0.0) _data[j](i+1) = get_smooth(vol,mask,coords[j],kernels);
      else _data[j](i+1) = vol(coords[j][0],coords[j][1],coords[j][2]);
    }
  }
  // Save coordinates
  _coords = coords;
} EddyCatch

float DataSelector::get_smooth(const NEWIMAGE::volume<float>&           ima,
			       const NEWIMAGE::volume<float>&           mask,
			       const std::vector<unsigned int>&         coords,
			       const std::vector<NEWMAT::ColumnVector>& kernels) EddyTry
{
  float oval = 0.0;
  float totwgt = 0.0;
  for (int k=0; k<kernels[2].Nrows(); k++) {
    float kwgt = kernels[2](k+1);
    for (int j=0; j<kernels[1].Nrows(); j++) {
      float jkwgt = kwgt*kernels[1](j+1);
      for (int i=0; i<kernels[0].Nrows(); i++) {
	int kk = coords[2] - kernels[2].Nrows()/2 + k;
	if (kk>=0 && kk<ima.zsize()) {
	  int jj = coords[1] - kernels[1].Nrows()/2 + j;
	  if (jj>=0 && jj<ima.ysize()) {
	    int ii = coords[0] - kernels[0].Nrows()/2 + i;
	    if (ii>=0 && ii<ima.xsize() && mask(ii,jj,kk)) {
	      float tmp = jkwgt*kernels[0](i+1);
	      totwgt += tmp;
	      oval += tmp*ima(ii,jj,kk);
	    }
	  }
	}
      }
    }
  }
  oval /= totwgt;
  return(oval);
} EddyCatch

double MMLHyParCF::cf(const NEWMAT::ColumnVector& p) const EddyTry
{
  std::vector<double> hp(p.Nrows());
  for (int i=0; i<p.Nrows(); i++) hp[i] = p(i+1);
  _K->SetHyperPar(hp);
  if (!_K->IsValid()) return(std::numeric_limits<double>::max());
  double ldKy = _K->LogDet();
  double neg_log_marg_ll = 0.0;
  if (_nthreads == 1) { // If we are to run single threaded
    for (unsigned int i=0; i<_data.size(); i++) neg_log_marg_ll += NEWMAT::DotProduct(_data[i],_K->iKy(_data[i]));
  }
  else { // We are to run multi-threaded
    std::vector<unsigned int> nvecs = EddyUtils::ScansPerThread(_data.size(),_nthreads);
    std::vector<std::thread> threads(_nthreads-1); // + main thread makes _nthreads
    std::vector<double> neg_log_marg_ll_vec(_nthreads,0.0);
    _K->CalculateInvK(); // Needed to allow for use of const _K->iKy for thread safety
    for (unsigned int i=0; i<_nthreads-1; i++) {
      threads[i] = std::thread(&MMLHyParCF::cf_helper,this,nvecs[i],nvecs[i+1],
			       std::ref(_data),_K,std::ref(neg_log_marg_ll_vec[i]));
    }
    this->cf_helper(nvecs[_nthreads-1],nvecs[_nthreads],_data,_K,neg_log_marg_ll_vec[_nthreads-1]);
    std::for_each(threads.begin(),threads.end(),std::mem_fn(&std::thread::join));
    neg_log_marg_ll = std::accumulate(neg_log_marg_ll_vec.begin(),neg_log_marg_ll_vec.end(),0.0);
  }
  neg_log_marg_ll += _data.size() * ldKy;
  neg_log_marg_ll *= 0.5;
  return(neg_log_marg_ll);
} EddyCatch

void MMLHyParCF::cf_helper(// Input
			   unsigned int first,
			   unsigned int last,
			   const std::vector<NEWMAT::ColumnVector>&  data,
			   const std::shared_ptr<const KMatrix>      K,
			   // Output
			   double&                                   neg_log_marg_ll) const EddyTry
{
  neg_log_marg_ll = 0.0;
  for (unsigned int i=first; i<last; i++) neg_log_marg_ll += NEWMAT::DotProduct(_data[i],_K->iKy(_data[i]));
  return;
} EddyCatch
			   
double CVHyParCF::cf(const NEWMAT::ColumnVector& p) const EddyTry
{
  std::vector<double> hp(p.Nrows());
  for (int i=0; i<p.Nrows(); i++) hp[i] = p(i+1);
  _K->SetHyperPar(hp);
  if (!_K->IsValid()) return(std::numeric_limits<double>::max());
  NEWMAT::ColumnVector ssd_vec(_K->NoOfScans());
  ssd_vec = 0.0;
  if (_nthreads == 1) { // If we are to run single threaded
    for (unsigned int i=0; i<_data.size(); i++) {
      NEWMAT::ColumnVector qn = _K->iKy(_data[i]);
      ssd_vec += NEWMAT::SP(qn,qn);
    }
  }
  else { // We are to run multi threaded
    std::vector<unsigned int> nvecs = EddyUtils::ScansPerThread(_data.size(),_nthreads);
    std::vector<std::thread> threads(_nthreads-1); // + main thread makes _nthreads
    std::vector<NEWMAT::ColumnVector> ssd_vec_vec(_nthreads,NEWMAT::ColumnVector(_K->NoOfScans()));
    _K->CalculateInvK(); // Needed to allow for use of const _K->iKy for thread safety
    for (unsigned int i=0; i<_nthreads-1; i++) {
      threads[i] = std::thread(&CVHyParCF::cf_helper,this,nvecs[i],nvecs[i+1],i,_nthreads,
			       std::ref(threads),std::ref(_data),_K,std::ref(ssd_vec_vec));
    }
    this->cf_helper(nvecs[_nthreads-1],nvecs[_nthreads],_nthreads-1,_nthreads,threads,_data,_K,ssd_vec_vec); // Takes care of joining
    ssd_vec = ssd_vec_vec[_nthreads-1];
  }
  NEWMAT::SymmetricMatrix iK = _K->iK();
  double ssd = 0.0;
  for (unsigned int i=0; i<_K->NoOfScans(); i++) ssd += ssd_vec(i+1) / sqr(iK(i+1,i+1));
  ssd /= _K->NoOfScans();
  return(ssd);
} EddyCatch

void CVHyParCF::cf_helper(// Input
			  unsigned int                              first,
			  unsigned int                              last,
			  unsigned int                              tid,       // Thread number
			  unsigned int                              numt,      // Number of threads
			  std::vector<std::thread>&                 threads,  
			  const std::vector<NEWMAT::ColumnVector>&  data,
			  const std::shared_ptr<const KMatrix>      K,
			  // Output
			  std::vector<NEWMAT::ColumnVector >&       ssd_vec_vec) const EddyTry
{
  // static std::mutex cout_mutex;

  // Do the multiplications and additions you are responsible for
  ssd_vec_vec[tid] = 0.0;
  for (unsigned int i=first; i<last; i++) {
    NEWMAT::ColumnVector qn = _K->iKy(_data[i]);
    ssd_vec_vec[tid] += NEWMAT::SP(qn,qn);
  }
  // Now add up results with other threads
  int step = 1;                                        // How far to the left your neighbour is
  while (step < numt) {
    if (is_odd(((int(numt)-1)-int(tid))/step)) return; // Is not a candidate
    int neighbour = int(tid)-step;                     // Check who your neighbour to the left is
    if (neighbour < 0) return;                         // Return if neighbour is out of bounds
    threads[neighbour].join();                         // Make sure that neighbour has finished
    // cout_mutex.lock();
    // std::cout << "I am thread # " << tid << ", and I will add thread # " << neighbour << " to myself" << std::endl;
    // cout_mutex.unlock();
    ssd_vec_vec[tid] += ssd_vec_vec[neighbour];
    step *= 2;
  }
  return; // At this stage the results whould have been summed up in ssd_vec_vec[numt-1]
} EddyCatch

double GPPHyParCF::cf(const NEWMAT::ColumnVector& p) const EddyTry
{
  std::vector<double> hp(p.Nrows());
  for (int i=0; i<p.Nrows(); i++) hp[i] = p(i+1);
  _K->SetHyperPar(hp);
  if (!_K->IsValid()) return(std::numeric_limits<double>::max());
  NEWMAT::ColumnVector gpp_vec(_K->NoOfScans());
  gpp_vec = 0.0;
  for (unsigned int i=0; i<_data.size(); i++) {
    NEWMAT::ColumnVector qn = _K->iKy(_data[i]);
    gpp_vec += NEWMAT::SP(qn,qn);
  }
  NEWMAT::SymmetricMatrix iK = _K->iK();
  for (unsigned int i=0; i<_K->NoOfScans(); i++) {
    gpp_vec(i+1) /= iK(i+1,i+1);
    gpp_vec(i+1) -= _data.size() * std::log(iK(i+1,i+1));
  }
  double gpp = 0.5 * gpp_vec.Sum() / _K->NoOfScans();
  return(gpp);
} EddyCatch

void CheapAndCheerfulHyParEstimator::Estimate(std::shared_ptr<const KMatrix>  K,
					      bool                            verbose) EddyTry
{
  std::shared_ptr<KMatrix> lK = K->Clone(); // Local copy of K
  lK->SetHyperPar(lK->GetHyperParGuess(_data));
  lK->MulErrVarBy(5.0); // Corresponds to ~ ff=5
  set_hpar(stl_2_newmat(lK->GetHyperPar()));
  if (verbose) cout << "Hyperparameters guesstimated to be: " << get_hpar() << endl;
} EddyCatch

void FullMontyHyParEstimator::Estimate(std::shared_ptr<const KMatrix>  K,
				       bool                            vbs) EddyTry
{
  std::shared_ptr<KMatrix> lK = K->Clone(); // Local copy of K
  set_hpar(stl_2_newmat(lK->GetHyperParGuess(_data)));
  if (_v) cout << "Initial guess for hyperparameters: " << get_hpar() << endl;
  _cf->SetData(_data);
  _cf->SetKMatrix(lK);
  MISCMATHS::NonlinParam nlpar(get_hpar().Nrows(),MISCMATHS::NL_NM,get_hpar());
  nlpar.SetVerbose(_v);
  nlpar.SetPrintWidthPrecision(10,5,15,10);
  nlpar.SetMaxIter(MITER);
  MISCMATHS::nonlin(nlpar,*_cf);
  set_hpar(nlpar.Par());
  if (_evff > 1.0) {
    if (_v) {
      cout << "Fudging parameters" << endl;
      cout << "Parameters start out as " << get_hpar() << endl;
    }
    std::shared_ptr<KMatrix> KK = lK->Clone();
    std::vector<double> hpar(get_hpar().Nrows());
    for (unsigned int i=0; i<hpar.size(); i++) hpar[i] = get_hpar()(i+1);
    KK->SetHyperPar(hpar);
    KK->MulErrVarBy(_evff);
    hpar = KK->GetHyperPar();
    NEWMAT::ColumnVector tmp_par(get_hpar().Nrows());
    for (unsigned int i=0; i<hpar.size(); i++) tmp_par(i+1) = hpar[i];
    set_hpar(tmp_par);
    if (_v) cout << "Parameters end up as " << get_hpar() << endl;
  }
  if (vbs) cout << "Estimated hyperparameters: " << get_hpar() << endl;
} EddyCatch