/*  AutoCorrEstimator.cc

    Mark Woolrich, FMRIB Image Analysis Group

    Copyright (C) 1999-2000 University of Oxford  */

/*  CCOPYRIGHT  */

#include <iostream>
#include <strstream>

#include "AutoCorrEstimator.h"
#include "miscmaths.h"
#include "Log.h"
#include "Volume.h"
#include "histogram.h"

using namespace NEWMAT;
using namespace UTILS;
namespace TACO {

  void AutoCorrEstimator::setDesignMatrix(const Matrix& dm) {
    Tracer tr("AutoCorrEstimator::setDesignMatrix");

    int sizeTS = xdata.getNumVolumes();
    int numPars = dm.Ncols();
    
    dminFFTReal.ReSize(zeropad, numPars);
    dminFFTImag.ReSize(zeropad, numPars);
      
    ColumnVector dmrow;
    dmrow.ReSize(zeropad);

    ColumnVector dm_fft_real, dm_fft_imag;
    ColumnVector dummy(zeropad);
    ColumnVector realifft(zeropad);

    // FFT design matrix
    for(int k = 1; k <= numPars; k++)
      {
	dummy = 0;
	dmrow = 0;
	mn(k) = MISCMATHS::mean(dm.Column(k));
	dmrow.Rows(1,sizeTS) = dm.Column(k) - mn(k);

	FFT(dmrow, dummy, dm_fft_real, dm_fft_imag);
	
	dminFFTImag.Column(k) = dm_fft_imag;
	dminFFTReal.Column(k) = dm_fft_real;
      } 
  }

  void AutoCorrEstimator::preWhiten(ColumnVector& in, ColumnVector& ret, int i, Matrix& dmret) {
      Tracer tr("AutoCorrEstimator::preWhiten");

      int sizeTS = xdata.getNumVolumes();
      int numPars = dminFFTReal.getNumSeries();

      ret.ReSize(sizeTS);
      dmret.ReSize(sizeTS, numPars);

      // FFT auto corr estimate
      dummy = 0;
      vrow = 0;
      vrow.Rows(1,sizeTS/2) = acEst.getSeries(i).Rows(1,sizeTS/2);
      vrow.Rows(zeropad - sizeTS/2 + 2, zeropad) = acEst.getSeries(i).Rows(2, sizeTS/2).Reverse();

      FFT(vrow, dummy, ac_fft_real, ac_fft_im);      

      // FFT x data
      dummy = 0;
      xrow = 0;
      xrow.Rows(1,sizeTS) = in;
     
      FFT(xrow, dummy, x_fft_real, x_fft_im);

      // inverse auto corr to give prewhitening filter:
      // no DC component so set first value to 0;
      ac_fft_real(1) = 0.0;	 
      for(int j = 2; j <= zeropad; j++)
	{
	  ac_fft_real(j) = 1.0/ac_fft_real(j);
	}

      // filter design matrix
      for(int k = 1; k <= numPars; k++)
	{
	  dm_fft_real = dminFFTReal.getSeries(k);
	  dm_fft_imag = dminFFTImag.getSeries(k);
	  FFTI(SP(ac_fft_real, dm_fft_real), SP(ac_fft_real, dm_fft_imag), realifft, dummy);

	  // place result into ret:
	  dmret.Column(k) = realifft.Rows(1,sizeTS);
	  float std = pow(MISCMATHS::var(dmret.Column(k)),0.5);
	  dmret.Column(k) = (dmret.Column(k)/std) + mn(k);
	}

      // Do filtering and inverse FFT:
      FFTI(SP(ac_fft_real, x_fft_real), SP(ac_fft_real, x_fft_im), realifft, dummy);
      
      // place result into ret:
      ret = realifft.Rows(1,sizeTS);     

    }

  void AutoCorrEstimator::preWhiten(VolumeSeries& in, VolumeSeries& ret)
    {
      Tracer tr("AutoCorrEstimator::preWhiten");

      cerr << "Prewhitening... ";

      int sizeTS = xdata.getNumVolumes();
      int numTS = xdata.getNumSeries();

      ret.ReSize(sizeTS, numTS);

      // make sure p_vrow is cyclic (even function)
      ColumnVector vrow, xrow;
      vrow.ReSize(zeropad);
      xrow.ReSize(zeropad);

      ColumnVector x_fft_real, ac_fft_real;
      ColumnVector x_fft_im, ac_fft_im;
      ColumnVector dummy(zeropad);    
      ColumnVector realifft(zeropad);
      int co = 1;

      for(int i = 1; i <= numTS; i++)
	{
	  // FFT auto corr estimate
	  dummy = 0;
	  vrow = 0;
	  vrow.Rows(1,sizeTS/2) = acEst.getSeries(i).Rows(1,sizeTS/2);
	  vrow.Rows(zeropad - sizeTS/2 + 2, zeropad) = acEst.getSeries(i).Rows(2, sizeTS/2).Reverse();
	  FFT(vrow, dummy, ac_fft_real, ac_fft_im);      
	  
	  // FFT x data
	  dummy = 0;
	  xrow = 0;
	  xrow.Rows(1,sizeTS/2) = in.getSeries(i).Rows(1,sizeTS/2);
	  xrow.Rows(zeropad - sizeTS/2 + 2, zeropad) = in.getSeries(i).Rows(2, sizeTS/2).Reverse();
	  FFT(xrow, dummy, x_fft_real, x_fft_im);

	  // inverse auto corr to give prewhitening filter:
	  // no DC component so set first value to 0;
	  ac_fft_real(1) = 0.0;	 
	  for(int j = 2; j <= zeropad; j++)
	    {
	      ac_fft_real(j) = 1.0/ac_fft_real(j);
	    }

	  // Do filtering and inverse FFT:
	  FFTI(SP(ac_fft_real, x_fft_real), SP(ac_fft_real, x_fft_im), realifft, dummy);
	
	  // place result into ret:
	  ret.Column(i) = realifft.Rows(1,sizeTS);

	  if(co > 100)
	    {
	      co = 1;
	      cerr << (float)i/(float)numTS << ",";
	    }
	  else
	    co++;
	}

      cerr << " Completed" << endl;
    }

  void AutoCorrEstimator::fitAutoRegressiveModel()
    {
      Tracer trace("AutoCorrEstimator::fitAutoRegressiveModel");
      
      cerr << "Fitting autoregressive model..." << endl;
      
      const int maxorder = 10;
      const int minorder = 0;

      int sizeTS = xdata.getNumVolumes();
      int numTS = xdata.getNumSeries();
      
      // setup temp variables
      ColumnVector x(sizeTS);
      ColumnVector order(numTS);
      ColumnVector betas(maxorder);
      acEst.ReSize(sizeTS, numTS);
      acEst = 0;
      int co = 1;
       
      for(int i = 1; i <= numTS; i++)
	{
	  x = xdata.getSeries(i).AsColumn();
	  order(i) = SIGPROC::Pacf(x, minorder, maxorder, betas);

	  if(order(i) != -1)
	    {
	      // Calculate auto corr:
	      ColumnVector Krow(sizeTS);
	      Krow = 0;
	      Krow(sizeTS) = 1;
	      if(order(i) > 0)
		Krow.Rows(sizeTS-order(i), sizeTS-1) = -betas.Rows(1,order(i)).Reverse();
	      
	      Matrix Kinv(sizeTS, sizeTS);
	      Kinv = 0;
	      for(int j = 1; j <= sizeTS; j++)
		{
		  Kinv.SubMatrix(j,j,1,j) = Krow.Rows(sizeTS-j+1,sizeTS).t();
		}
	      
	      // Kinv now becomes V:
	      Kinv = (Kinv.t()*Kinv).i();
	      acEst.SubMatrix(1,sizeTS/2+1,i,i) = (Kinv.SubMatrix(sizeTS/2, sizeTS/2, sizeTS/2, sizeTS)/Kinv.MaximumAbsoluteValue()).AsColumn();
	       
	      if(co > 200)
		{
		   co = 1;
		   cerr << (float)i/(float)numTS << ",";
		 }
	      else
		co++;
	    }
	}
      
      Log::getInstance().out("order", order);
      cerr << " Completed" << endl; 
    }
  
  int AutoCorrEstimator::establishUsanThresh(const Volume& epivol)
    {
      int usanthresh = 100;      
      int num = epivol.getVolumeSize();

      Histogram hist(epivol, num/200);
      hist.generate();
      float mode = hist.mode();
      cerr << "mode = " << mode << endl;
 
      float sum = 0.0;
      int count = 0;

      // Work out standard deviation from mode for values greater than mode:
      for(int i = 1; i <= num; i++) {
	if(epivol(i) > mode) {
	  sum = sum + (epivol(i) - mode)*(epivol(i) - mode);
	  count++;
	}
      }

      int sig = (int)pow(sum/num, 0.5);
      cerr << "sig = " << sig << endl;

      usanthresh = sig/3;

      return usanthresh;
    } 

  void AutoCorrEstimator::spatiallySmooth(const string& usanfname, const Volume& epivol, int masksize, const string& epifname, const string& susanpath) {
    Tracer trace("AutoCorrEstimator::spatiallySmooth");
    
    Log& logger = Log::getInstance();

    // Establish epi thresh to use:
    int usanthresh = establishUsanThresh(epivol);
     
    // Setup external call to susan program:
    char callsusanstr[1000];
    ostrstream osc3(callsusanstr,1000);
    string preSmoothVol = "preSmoothVol";
    string postSmoothVol = "postSmoothVol";

    osc3 << susanpath << " "
	 << logger.getDir() << "/" << preSmoothVol << " 1 " 
	 << logger.getDir() << "/" << postSmoothVol << " " 
	 << masksize << " 3D 0 1 " << usanfname << " " << usanthresh << " " 
	 << logger.getDir() << "/" << "usanSize" << '\0';
    
    // Loop through first third of volumes
    // assume the rest are zero
    int factor = 10000;

    // Setup volume for reading and writing volumes:
    Volume vol(acEst.getNumSeries(), xdata.getDims(), xdata.getPreThresholdPositions());
   
    int i = 2;
    
    cerr << "Spatially smoothing auto corr estimates" << endl;

    cerr << callsusanstr << endl;

    for(; i < 20; i++)
      {
	// output unsmoothed estimates:
	vol = acEst.getVolume(i).AsColumn()*factor;
	vol.unthreshold();
	vol.writeAsInt(logger.getDir() + "/" + preSmoothVol);

	// call susan:
	system(callsusanstr);

	// read in smoothed volume:
	vol.read(logger.getDir() + "/" + postSmoothVol);
	vol.threshold();
	acEst.getVolume(i) = vol.AsRow()/factor;
	cerr << ".";
      }
    cerr << endl;

    // Clear unwanted written files
    char rmfilesstr[1000];
    ostrstream osc(rmfilesstr,1000);
    osc << "rm -rf " 
	<< logger.getDir() + "/" + postSmoothVol + "* "
	<< logger.getDir() + "/" + preSmoothVol + "* "
	<< logger.getDir() + "/" + epifname + "* "
	<< logger.getDir() + "/usanSize*" << '\0';
    
    cerr << rmfilesstr << endl;

    system(rmfilesstr);

    cerr << "Completed" << endl;
  }

  void AutoCorrEstimator::calcRaw() { 
    
    cerr << "Calculating raw AutoCorrs...";      

    AutoCorr(xdata, acEst, zeropad);

    cerr << " Completed" << endl;  
  }
  
  void AutoCorrEstimator::filter(const ColumnVector& filterFFT) {

    Tracer tr("AutoCorrEstimator::filter");

    cerr << "Combining temporal filtering effects with AutoCorr estimates... ";

    // This function adjusts the autocorrelations as if the
    // xdata has been filtered by the passed in filterFFT
    // DOES NOT filter the xdata itself
    ColumnVector vrow;

    // make sure p_vrow is cyclic (even function)
    vrow.ReSize(zeropad);

    ColumnVector fft_real;
    ColumnVector fft_im;
    ColumnVector dummy(zeropad);    
    ColumnVector realifft(zeropad);
    int sizeTS = xdata.getNumVolumes();

    for(int i = 1; i <= xdata.getNumSeries(); i++)
      {
	dummy = 0;
	vrow = 0;
	vrow.Rows(1,sizeTS/2) = acEst.getSeries(i).Rows(1,sizeTS/2);
	vrow.Rows(zeropad - sizeTS/2 + 2, zeropad) = acEst.getSeries(i).Rows(2, sizeTS/2).Reverse();
      
	FFT(vrow, dummy, fft_real, fft_im);

	FFTI(SP(fft_real, filterFFT), dummy, realifft, dummy);
	
	// place result into acEst:
	acEst.Column(i) = realifft.Rows(1,sizeTS)/realifft(1);
      }

    cerr << " Completed" << endl;
    
  }
  
  void AutoCorrEstimator::pava() {
    
    Tracer tr("AutoCorrEstimator::pava");
    
    cerr << "Using New PAVA on AutoCorr estimates... ";

    for(int i = 1; i <= xdata.getNumSeries(); i++) {
	int sizeTS = xdata.getNumVolumes();
	int stopat = (int)sizeTS/2;

	// 5% point of distribution of autocorr about zero
	const float th = (-1/sizeTS)+(2/sqrt(sizeTS));

	ColumnVector values = acEst.Column(i);
	ColumnVector zero(1);
	zero = 0;
	values = values.Rows(1,stopat) & zero;

	ColumnVector gm(stopat + 1);
	for(int j = 1; j <= stopat + 1; gm(j) = j++);
	  
	ColumnVector weights(stopat+1);
	weights = 1;

	bool anyviolators = true;

	while(anyviolators) {
	  anyviolators = false;
	  
	  for(int k = 2; k <= values.Nrows(); k++) {
	    if(values(k) > values(k-1)) {
	      anyviolators = true;
	      values(k-1) = (values(k-1)*weights(k-1) + values(k)*weights(k))/(weights(k-1) + weights(k));
	      values = values.Rows(1,k-1) & values.Rows(k+1,values.Nrows());
	      weights(k-1) = weights(k) + weights(k-1);
	      weights = weights.Rows(1,k-1) & weights.Rows(k+1,weights.Nrows());

	      for(int j = 1; j <= stopat + 1; j++) {
		if(gm(j) >= k)  
		  gm(j) = gm(j)-1;
	      }

	      break;
	    }
	  }
	}
	
	acEst.Column(i) = 0.0;
	int j=1;

	for(; j <= stopat; j++) {
	  
	  acEst(j,i) = values(gm(j));
	  if(acEst(j,i) <= 0.0)
	    {
	      acEst(j,i) = 0.0;
	      break;
 	    }
	}
	
	if(acEst(2,i) < th/2)
	{
	acEst.SubMatrix(2,stopat,i,i) = 0;
	}

	else if(j > 2)
	  //if(j > 2)  
	{
	    int endst = j;
	    int stst = j-(int)(1+(j/8.0));

	    const int expwidth = MISCMATHS::Max((endst - stst)/2,1);
	    const int exppow = 2;
	
	    for(j = stst; j <= endst; j++)
	      {
		acEst(j,i) = acEst(j,i)*exp(-pow((j-stst)/expwidth,exppow));
	      }	
	}
	
    }
    cerr << " Completed" << endl;
  }
  
  void AutoCorrEstimator::applyConstraints() {

    Tracer tr("AutoCorrEstimator::applyConstraints");

    cerr << "Applying constraints to AutoCorr estimates... ";

    for(int i = 1; i <= xdata.getNumSeries(); i++)
      {
	int sizeTS = xdata.getNumVolumes();
	int j = 3;
	int stopat = (int)sizeTS/4;

	// found1 is last valid value above threshold
	int found1 = stopat;

	// 5% point of distribution of autocorr about zero
	const float thresh = (-1/sizeTS)+(2/sqrt(sizeTS));

	acEst(2,i) = (acEst(2,i)+ acEst(3,i))/2;
	if(acEst(2,i) < 0)
	  {
	    acEst(2,i) = 0;
	  }
	else
	  {
	    float grad = 0.0;

	    while(j <= stopat && j < found1 + 2)
	      {    
		grad = ((acEst(j,i) + acEst(j+1,i))/2 - acEst(j-1,i))/1.5;
		if(grad < 0)
		  acEst(j,i) = grad + acEst(j-1,i);
		else
		  acEst(j,i) = acEst(j-1,i);
		
		// look for threshold
		if(acEst(j,i) < thresh/3.0 && found1 == stopat)
		  {
		    found1 = j;
		  }
	   
		if(acEst(j,i) < 0)
		  {
		    acEst(j,i) = 0;
		  }
		
		j++;
	      }
	  }

	// set rest to zero:
	for(; j <= sizeTS; j++)
	  {
	      acEst(j,i) = 0;
	  }
      }
    cerr << "Completed" << endl;
  }

  void AutoCorrEstimator::getMeanEstimate(ColumnVector& ret)
    {
      Tracer tr("AutoCorrEstimator::getMeanEstimate");

      ret.ReSize(acEst.getNumVolumes());
      // Calc global Vrow:
      for(int i = 1; i <= acEst.getNumVolumes(); i++)
	{
	  ret(i) = MISCMATHS::mean(acEst.getVolume(i).AsColumn());
	}
    }
}