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

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

#include <cstdlib>
#include <string>
#include <vector>
#include <cmath>
#pragma push
#pragma diag_suppress = code_is_unreachable // Supress warnings from armawrap
#pragma diag_suppress = expr_has_no_effect  // Supress warnings from boost
#include "armawrap/newmat.h"
#include "newimage/newimageall.h"
#pragma pop
#include "miscmaths/miscmaths.h"
#include "EddyHelperClasses.h"
#include "EddyUtils.h"
#include "fmriPredictor.h"
#include "CudaVolume.h"

using namespace EDDY;

/****************************************************************//**
* \brief Returns prediction for point given by indx
*
* Returns a predicted image for the given index, where index refers
* to the corresponding time-point.
* It should be noted that this implementation isn't very efficient as
* it reloads all image to the GPU. That means that if there are N images
* and we want to predict them all it takes N*N transfers to the GPU.
* \param index refers to the corresponding time-point
* \param exclude Decides if indx itself should be used in the prediction
* (exclude=false) or not (exclude=true)
* \param pvec Prediction vector for index
* \param pi The "predicted image"
*
********************************************************************/
void fmriPredictor::predict_image_gpu(// Input
				      unsigned int             indx,
				      bool                     exclude,
				      const arma::rowvec&      pvec,
				      // Output
				      NEWIMAGE::volume<float>& pi) const EddyTry
{
  unsigned int pvec_length = (exclude) ? this->no_of_scans_in_same_session(indx)-1 : this->no_of_scans_in_same_session(indx);
  if (pvec_length != pvec.n_cols) throw EddyException("fmriPredictor::predict_image_gpu: Wrong length pvec");

  if (!NEWIMAGE::samesize(pi,*(this->first_imptr()))) {
    pi.reinitialize(this->first_imptr()->xsize(),this->first_imptr()->ysize(),this->first_imptr()->zsize());
    NEWIMAGE::copybasicproperties(*(this->first_imptr()),pi);
  }
  EDDY::CudaVolume pcv(pi,false);
  int sindx = this->index_in_session(indx);
  std::shared_ptr<NEWIMAGE::volume<float> > ptr=nullptr;
  for (unsigned int s=0; s<this->no_of_scans_in_same_session(indx); s++) {
    if (exclude) {
      if (s < sindx) pcv.MultiplyAndAddToMe(EDDY::CudaVolume(*(this->imptr_from_same_session(indx,s))),pvec(s));
      // Do nothing if s==sindx
      else if (s > indx) pcv.MultiplyAndAddToMe(EDDY::CudaVolume(*(this->imptr_from_same_session(indx,s))),pvec(s-1));
    }
    else pcv.MultiplyAndAddToMe(EDDY::CudaVolume(*(this->imptr_from_same_session(indx,s))),pvec(s));
  }
  pcv += EDDY::CudaVolume(*(this->meanptr_from_same_session(indx)));
  pcv.GetVolume(pi);
  return;
} EddyCatch

void fmriPredictor::predict_images_gpu(// Input
				       const std::vector<unsigned int>&       indicies,
				       bool                                   exclude,
				       const std::vector<arma::rowvec>&       pvecs,
				       // Output
				       std::vector<NEWIMAGE::volume<float> >& pi) const EddyTry
{
  // Sanity checks
  if (indicies.size() != pvecs.size() || indicies.size() != pi.size()) {
    throw EDDY::EddyException("fmriPredictor::predict_images_gpu: mismatch among indicies, pvecs and pi");
  }
  if (!this->all_in_same_session(indicies)) throw EDDY::EddyException("fmriPredictor::predict_images_gpu: indicies are in different sessions");
  unsigned int pvec_length = (exclude) ? this->no_of_scans_in_same_session(indicies[0])-1 : this->no_of_scans_in_same_session(indicies[0]);
  if (std::any_of(pvecs.begin(),pvecs.end(),[pvec_length](const arma::rowvec& rv){ return(rv.n_cols != pvec_length); })) {
    throw EDDY::EddyException("fmriPredictor::predict_images_gpu: pvec with wrong length");
  }

  // Start by allocating space on GPU for all output images
  std::vector<EDDY::CudaVolume> pcvs(indicies.size());
  for (unsigned int i=0; i<indicies.size(); i++) {
    if (!NEWIMAGE::samesize(pi[i],*(this->first_imptr()))) {
      pi[i].reinitialize(this->first_imptr()->xsize(),this->first_imptr()->ysize(),this->first_imptr()->zsize());
      NEWIMAGE::copybasicproperties(*(this->first_imptr()),pi[i]);
    }
    pcvs[i].SetHdr(pi[i]);
  }
  // Transfer mean image to the GPU
  EDDY::CudaVolume mean = *(this->meanptr_from_same_session(indicies[0]));
  // Do the GP predictions
  for (unsigned int s=0; s<this->no_of_scans_in_same_session(indicies[0]); s++) { // s index into original volumes
    EDDY::CudaVolume cv = *(this->imptr_from_same_session(indicies[0],s));
    for (unsigned int i=0; i<indicies.size(); i++) { // i index into predictions
      if (exclude) {
	if (s < this->index_in_session(indicies[i])) pcvs[i].MultiplyAndAddToMe(cv,pvecs[i](s));
	// Do nothing if (s == indicies[i])
	else if (s > this->index_in_session(indicies[i])) pcvs[i].MultiplyAndAddToMe(cv,pvecs[i](s-1));
      }
      else pcvs[i].MultiplyAndAddToMe(cv,pvecs[i](s));
    }
  }
  // Add means to predictions and transfer back from GPU
  for (unsigned int i=0; i<indicies.size(); i++) {
    pcvs[i] += mean;
    pcvs[i].GetVolume(pi[i]);
  }
  return;
} EddyCatch