/*! \file fmriPredictor.cpp
    \brief Contains definitions for class for making Gaussian process based predictions about fMRI data.

    \author Jesper Andersson
    \version 1.0b, Feb., 2022.
*/
// Definitions of class to make Gaussian-Process
// based predictions about diffusion data.
//
// fmriPredictor.cpp
//
// Jesper Andersson, FMRIB Image Analysis Group
//
// Copyright (C) 2022 University of Oxford
//

#include <cstdlib>
#include <string>
#include <vector>
#include <cmath>
#include "armawrap/newmat.h"
#include "newimage/newimageall.h"
#include "miscmaths/miscmaths.h"
#include "EddyHelperClasses.h"
#include "EddyUtils.h"
#include "KMatrix.h"
#include "HyParEstimator.h"
#include "fmriPredictor.h"

using namespace EDDY;

NEWIMAGE::volume<float> fmriPredictor::Predict(unsigned int indx,
					       bool         exclude) const EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::Predict:const: Not yet fully populated");
  if (!IsValid()) throw EddyException("fmriPredictor::Predict:const: Not yet ready for predictions");
  arma::rowvec pv = _Kmats[_glist[indx]._sess]->PredVec(_glist[indx]._sindx,exclude); // Calls const version
  NEWIMAGE::volume<float> pi = *(_slist[0].SPtr(0)); pi = 0.0;
  #ifdef COMPILE_GPU
  predict_image_gpu(indx,exclude,pv,pi);
  #else
  predict_image_cpu(indx,exclude,pv,pi);
  #endif
  return(pi);
} EddyCatch

NEWIMAGE::volume<float> fmriPredictor::Predict(unsigned int indx,
					       bool         exclude) EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::Predict:non-const: Not yet fully populated");
  if (!IsValid()) throw EddyException("fmriPredictor::Predict:non-const: Not yet ready for predictions");
  arma::rowvec pv = _Kmats[_glist[indx]._sess]->PredVec(_glist[indx]._sindx,exclude); // Calls non-const version
  NEWIMAGE::volume<float> pi = *(_slist[0].SPtr(0)); pi = 0.0;
  #ifdef COMPILE_GPU
  predict_image_gpu(indx,exclude,pv,pi);
  #else
  predict_image_cpu(indx,exclude,pv,pi);
  #endif
  return(pi);
} EddyCatch

NEWIMAGE::volume<float> fmriPredictor::PredictCPU(unsigned int indx,
						  bool         exclude) EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::Predict:non-const: Not yet fully populated");
  if (!IsValid()) throw EddyException("fmriPredictor::Predict:non-const: Not yet ready for predictions");
  arma::rowvec pv = _Kmats[_glist[indx]._sess]->PredVec(_glist[indx]._sindx,exclude); // Calls non-const version
  NEWIMAGE::volume<float> pi = *(_slist[0].SPtr(0)); pi = 0.0;
  predict_image_cpu(indx,exclude,pv,pi);
  return(pi);
} EddyCatch

std::vector<NEWIMAGE::volume<float> > fmriPredictor::Predict(const std::vector<unsigned int>& indicies,
							     bool                             exclude) EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::Predict: Not yet fully populated");
  if (!IsValid()) throw EddyException("fmriPredictor::Predict: Not yet ready for predictions");
  std::vector<NEWIMAGE::volume<float> > pi(indicies.size());
  std::vector<arma::rowvec> pvecs(indicies.size());
  for (unsigned int i=0; i<indicies.size(); i++) pvecs[i] = _Kmats[_glist[indicies[i]]._sess]->PredVec(_glist[indicies[i]]._sindx,exclude);
  #ifdef COMPILE_GPU
  predict_images_gpu(indicies,exclude,pvecs,pi);
  #else
  predict_images_cpu(indicies,exclude,pvecs,pi);
  #endif
  return(pi);
} EddyCatch

NEWIMAGE::volume<float> fmriPredictor::InputData(unsigned int indx) const EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::InputData: Not yet fully populated");
  if (indx >= _glist.size()) throw EddyException("DiffusionGP::InputData: indx out of range");
  return(*(_slist[_glist[indx]._sess].SPtr(_glist[indx]._sindx)) + *(_mptrs[_glist[indx]._sess]));
} EddyCatch

std::vector<NEWIMAGE::volume<float> > fmriPredictor::InputData(const std::vector<unsigned int>& indx) const EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::InputData: Not yet fully populated");
  for (unsigned int i=0; i<indx.size(); i++) if (indx[i] >= _glist.size()) throw EddyException("fmriPredictor::InputData: indx out of range");
  std::vector<NEWIMAGE::volume<float> > rval(indx.size());
  for (unsigned int i=0; i<indx.size(); i++) rval[i] = *(_slist[_glist[indx[i]]._sess].SPtr(_glist[indx[i]]._sindx)) + *(_mptrs[_glist[indx[i]]._sess]);
  return(rval);
} EddyCatch

double fmriPredictor::PredictionVariance(unsigned int indx,
					 bool         exclude) EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::PredictionVariance:const: Not yet fully populated");
  if (!IsValid()) throw EddyException("fmriPredictor::PredictionVariance:const: Not yet ready for predictions");
  double pv = _Kmats[_glist[indx]._sess]->PredVar(_glist[indx]._sindx,exclude);
  return(pv);
} EddyCatch

double fmriPredictor::ErrorVariance(unsigned int indx) const EddyTry
{
  if (!IsPopulated()) throw EddyException("fmriPredictor::ErrorVariance:const: Not yet fully populated");
  if (!IsValid()) throw EddyException("fmriPredictor::ErrorVariance:const: Not yet ready for predictions");
  double ev = _Kmats[_glist[indx]._sess]->ErrVar(_glist[indx]._sindx);
  return(ev);
} EddyCatch

void fmriPredictor::SetNoOfScans(unsigned int n) EddyTry
{
  if (n == _glist.size()) return; // No change
  else if (n > _glist.size()) {   // If increasing size
    std::lock_guard<std::mutex>  lg(_set_mut);
    _glist.resize(n); // New elements populated according to default constructor
    _lak = false;
    for (unsigned int i; i<_slist.size(); i++) _Kmats[i]->Reset();
  }
  else { // Decreasing size not allowed
    throw EddyException("fmriPredictor::SetNoOfScans: Decreasing size not allowed");
  }
  return;
} EddyCatch

void fmriPredictor::SetScan(const NEWIMAGE::volume<float>& scan, 
			    const DiffPara&                dp,
			    unsigned int                   indx,
			    unsigned int                   sess) EddyTry
{
  std::lock_guard<std::mutex> lg(_set_mut);
  // First just sanity checks
  if (indx >= _glist.size()) throw EddyException("fmriPredictor::SetScan: Invalid image index"); 
  if (_tr < 0.0) _tr = dp.TR();
  else if (std::abs(_tr - dp.TR()) > 1e-6) throw EddyException("fmriPredictor::SetScan: You cannot mix scans with different repetition times"); 
  if (_glist.size() && !NEWIMAGE::samesize(*_slist[_glist[0]._sess].SPtr(0),scan)) throw EddyException("fmriPredictor::SetScan: Wrong image dimension");
  if (_glist[indx]._sess >= 0 && _glist[indx]._sess != sess) throw EddyException("fmriPredictor::SetScan: You cannot change session of a scan");

  if (_glist[indx]._sess >= 0 && _slist[_glist[indx]._sess].Gndx(_glist[indx]._sindx) != static_cast<int>(indx)) throw EddyException("fmriPredictor::SetScan: Lists are inconsistent");
  // Next we do the job
  if (_glist[indx]._sess >= 0) { // If the same index has already been loaded before
    _slist[_glist[indx]._sess].SPtr(_glist[indx]._sindx) = std::make_shared<NEWIMAGE::volume<float> >(scan);
  }
  else { // If this is the first time this index is loaded
    _glist[indx]._sess = sess;
    _glist[indx]._sindx = _slist[sess].Size();
    _slist[sess].PushBack(scan,indx);
  }
  _lak = false;
} EddyCatch

void fmriPredictor::EvaluateModel(const NEWIMAGE::volume<float>& mask, float fwhm, bool verbose) EddyTry
{
  // The first thing we want to do is to "clean" up the lists.
  if (!lists_are_kosher()) throw EddyException("fmriPredictor::lists_are_kosher: Lists are inconsistent");
  // Next make one K-matrix per session
  _Kmats.resize(_slist.size());
  for (unsigned int i=1; i<_slist.size(); i++) {
    _Kmats[i] = _Kmats[0]->Clone();
  }
  // Populate K-matrices with distances. Note that we have shoe-horned it into a "diffusion syntax"
  for (unsigned int i=0; i<_slist.size(); i++) {
    std::vector<DiffPara> dpars(_slist[i].Size(),DiffPara(-100,_tr));
    _Kmats[i]->SetDiffusionPar(dpars);
  }
  mean_correct(); // Mean correct (on a per session basis)
  // Next we need to select data for the estimation of hyper-parameters.
  // If all sessions have the same number of scans we can mix-and-match data
  // from different sessions. If not we pick a session at random, but with 
  // probability weighted by the number of scans in the session.
  if (same_no_of_scans_in_all_sessions()) { // Mix-and-match
    DataSelector ds = DataSelector(_slist[0].SPtr_list(),mask,(_hpe->GetNVox()/_slist.size())+1,FourthDimension(0),fwhm,_hpe->RndInit());
    for (unsigned int i=1; i<_slist.size(); i++) {
      ds.ConcatToMe(DataSelector(_slist[i].SPtr_list(),mask,(_hpe->GetNVox()/_slist.size())+1,FourthDimension(i),fwhm,_hpe->RndInit()));
    }
    _hpe->SetData(ds.GetData());
    _hpe->Estimate(_Kmats[0],verbose);
    for (unsigned int i=0; i<_slist.size(); i++) {
      _Kmats[i]->SetHyperPar(_hpe->GetHyperParameters());
      _Kmats[i]->CalculateInvK();
    }    
  }
  else { // Pick random session
    // Draw random number 0--no_of_scans-1
    int seed = (_hpe->RndInit()) ? _hpe->RndInit() : static_cast<int>(time(NULL));
    std::mt19937 gen(seed);
    std::uniform_int_distribution<> distr(0,static_cast<int>(no_of_scans())-1);
    int rndnr = distr(gen);
    // Determine which session this corresponds to
    int cumscans = 0;
    unsigned int si;
    for (si=0; si<_slist.size(); si++) {
      cumscans += _slist[si].Size();
      if (cumscans >= rndnr) break;
    }
    DataSelector ds = DataSelector(_slist[si].SPtr_list(),mask,_hpe->GetNVox(),FourthDimension(si),fwhm,_hpe->RndInit()); 
    _hpe->SetData(ds.GetData());
    _hpe->Estimate(_Kmats[si],verbose);
    for (unsigned int i=0; i<_slist.size(); i++) {
      _Kmats[i]->SetHyperPar(_hpe->GetHyperParameters());
      _Kmats[i]->CalculateInvK();
    }
  }
  return;
} EddyCatch

void fmriPredictor::WriteImageData(const std::string& fname) const EddyTry
{
  char ofname[256];
  if (!IsPopulated() || !_lak) throw EddyException("fmriPredictor::WriteImageData: Not yet fully populated");
  // For practical reasons the volumes are written individually
  for (unsigned int i=0; i<_glist.size(); i++) {
    sprintf(ofname,"%s_%03d_%02d_%03d",fname.c_str(),i,_glist[i]._sess,_glist[i]._sindx);
    NEWIMAGE::write_volume(*(_slist[_glist[i]._sess].SPtr(_glist[i]._sindx)),ofname);
  }
  for (unsigned int i=0; i<_mptrs.size(); i++) {
    sprintf(ofname,"%s_mean_%02d",fname.c_str(),i);
    NEWIMAGE::write_volume(*(_mptrs[i]),ofname);
  }
} EddyCatch

void fmriPredictor::WriteMetaData(const std::string& fname) const EddyTry
{
  if (!IsPopulated() || !_lak) throw EddyException("fmriPredictor::WriteMetaData: Not yet fully populated");
  char ofname[256];
  for (unsigned int i=0; i<_slist.size(); i++) {
    sprintf(ofname,"%s_%02d",fname.c_str(),i);
    _Kmats[i]->Write(ofname);
  }
} EddyCatch

void fmriPredictor::mean_correct() EddyTry
{
  NEWIMAGE::volume<float> mean = *_slist[_glist[0]._sess].SPtr(_glist[0]._sindx);
  for (unsigned int li=0; li<_slist.size(); li++) {
    mean = *_slist[li].SPtr(0);
    for (unsigned int i=1; i<_slist[li].Size(); i++) mean += *_slist[li].SPtr(i);
    mean /= static_cast<float>(_slist[li].Size());
    for (unsigned int i=0; i<_slist[li].Size(); i++) *_slist[li].SPtr(i) -= mean;
  }
} EddyCatch

bool fmriPredictor::is_populated() const EddyTry
{
  for (unsigned int i=0; i<_glist.size(); i++) if (_glist[i]._sess < 0) return(false);
  return(true);
} EddyCatch

bool fmriPredictor::same_no_of_scans_in_all_sessions() const EddyTry
{
  for (unsigned int i=1; i<_slist.size(); i++) if (_slist[i].Size() != _slist[0].Size()) return(false);
  return(true);
} EddyCatch

/****************************************************************//**
*
* This ensures that the entries in the _slist are sorted in ascending 
* order of "global index". But when reorganising the _slist we also 
* need to make sure that the entries in the _glist still points to 
* the right _slist entries.
* 
********************************************************************/
bool fmriPredictor::lists_are_kosher() EddyTry
{
  if (!_lak) {
    // First sort the entries in the _slists with the _gindx as key
    for (unsigned int li=0; li<_slist.size(); li++) _slist[li].SortByGndx();
    // Then make sure that _glist and _slist are consistent
    for (unsigned int li=0; li<_slist.size(); li++) {
      for (unsigned int i=0; i<_slist[li].Size(); i++) {
	if (_glist[_slist[li].Gndx(i)]._sess != li) return(false); 
	_glist[_slist[li].Gndx(i)]._sindx = i;
      }
    }
    // Then do a little sanity check of _glist
    for (unsigned int i=0; i<_glist.size(); i++) {
      if (_glist[i]._sess >= _slist.size() || _glist[i]._sindx >= _slist[_glist[i]._sess].Size()) return(false);
    }
    // And finally check that _glist don't have duplicates
    if (glist_has_duplicates()) return(false);
  }
  return(true);
} EddyCatch

bool fmriPredictor::glist_has_duplicates() const EddyTry
{
  for (unsigned int i=0; i<_glist.size()-1; i++) {
    for (unsigned int j=i+1; j<_glist.size(); j++) {
      if (_glist[i]._sess == _glist[j]._sess && _glist[i]._sindx == _glist[j]._sindx) return(true);
    }
  }
  return(false);
} EddyCatch

void fmriPredictor::predict_image_cpu(// Input
				      unsigned int             indx,
				      bool                     exclude,
				      const arma::rowvec&      pv,
				      // Output
				      NEWIMAGE::volume<float>& pi) const EddyTry
{
  unsigned int ys = (exclude) ? _slist[_glist[indx]._sess].Size()-1 : _slist[_glist[indx]._sess].Size();
  arma::colvec y(ys);
  pi = *_mptrs[_glist[indx]._sess];
  for (int k=0; k<pi.zsize(); k++) {
    for (int j=0; j<pi.ysize(); j++) {
      for (int i=0; i<pi.xsize(); i++) {
	if (get_y(i,j,k,indx,exclude,y)) pi(i,j,k) += static_cast<float>(arma::as_scalar(pv*y));
	else pi(i,j,k) = 0.0;
      }
    }
  }
} EddyCatch

void fmriPredictor::predict_images_cpu(// Input
				       const std::vector<unsigned int>&       indicies,
				       bool                                   exclude,
				       const std::vector<arma::rowvec>&       pvecs,
				       // Output
				       std::vector<NEWIMAGE::volume<float> >& pi) const EddyTry
{
  if (indicies.size() != pvecs.size() || indicies.size() != pi.size()) {
    throw EDDY::EddyException("fmriPredictor::predict_images_cpu: mismatch among indicies, pvecs and pi");
  }
  for (unsigned int i=0; i<indicies.size(); i++) predict_image_cpu(indicies[i],exclude,pvecs[i],pi[i]);
  return;
} EddyCatch

bool fmriPredictor::get_y(// Input
			  unsigned int i, unsigned int j, unsigned int k, unsigned int indx, bool exclude,
			  // Output
			  arma::colvec&  y) const EddyTry
{
  const std::vector<std::shared_ptr<NEWIMAGE::volume<float> > >& sptrl = _slist[_glist[indx]._sess].SPtr_list();
  unsigned int sindx = _glist[indx]._sindx;
  for (unsigned int t=0, tt=0; t<sptrl.size(); t++) {
    if (!exclude || t!=sindx) {
      if (!(*sptrl[t])(i,j,k)) return(false);
      else y(tt++) = (*sptrl[t])(i,j,k);
    }
  }
  return(true);
} EddyCatch

void fmriPredictor::ptr_index_list::SortByGndx() EddyTry
{
  std::vector<int> indicies(this->Size());
  std::iota(indicies.begin(),indicies.end(),0);
  std::sort(indicies.begin(),indicies.end(),
	    [this](int a, int b){ return(this->Gndx(a) < this->Gndx(b)); });
  std::vector<std::shared_ptr<NEWIMAGE::volume<float> > > sptr_lst(this->Size());
  std::vector<int> gndx_lst(this->Size());
  for (unsigned int i=0; i<this->Size(); i++) {
    sptr_lst[i] = this->_sptr_lst[indicies[i]];
    gndx_lst[i] = this->_gndx_lst[indicies[i]];
  }
  this->_sptr_lst = sptr_lst;
  this->_gndx_lst = gndx_lst;
} EddyCatch