/*  AutoCorrEstimator.cc

    Mark Woolrich, FMRIB Image Analysis Group

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

/*  CCOPYRIGHT  */

#include <iostream>
#include <strstream>
#define WANT_STREAM

#include "AutoCorrEstimator.h"
#include "miscmaths.h"
#include "Log.h"
#include "Volume.h"
#include "histogram.h"
#include "miscmaths.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(ColumnVector(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, bool highfreqremovalonly) {
      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);      

      float norm = ac_fft_real.SumSquare();
      
      // Compare with raw FFT to detect high frequency artefacts:
      bool violate = false;
      ColumnVector violators(zeropad);
      violators = 1;
      for(int j = 1; j <= zeropad; j++)
	{
	  E(j,i) = sqrt(E(j,i)/((ac_fft_real(j)*ac_fft_real(j))/norm));

	  // look for high frequency artefacts
	  if(E(j,i) > 4 && j > zeropad/4 && j < 3*zeropad/4)
	    {	  
	      violate = true;
	      violators(j) = 0;
	      countLargeE(j) = countLargeE(j) + 1;
	    }
	}

      // FFT x data
      dummy = 0;
      xrow = 0;
      xrow.Rows(1,sizeTS) = in;
     
      FFT(xrow, dummy, x_fft_real, x_fft_im);
      
      if(highfreqremovalonly)
	{
	  ac_fft_real = violators;
	}
      else
	{
	  // 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/sqrt(fabs(ac_fft_real(j)));	      
	    }
	}

      // normalise ac_fft such that sum(j)(ac_fft_real)^2 = 1
      ac_fft_real /= sqrt(ac_fft_real.SumSquare()/zeropad);

      // 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) + mn(k);
	  //float std = pow(MISCMATHS::var(ColumnVector(dmret.Column(k))),0.5);
	  //dmret.Column(k) = (dmret.Column(k)/std) + mn(k);
	}

      // Do filtering of data:
      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);
	    }

	  // normalise ac_fft such that sum(j)(ac_fft_real)^2 = 1
	  ac_fft_real /= sqrt(ac_fft_real.SumSquare()/zeropad);
	  
	  // 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 = 15;
      const int minorder = 1;

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

	  order(i) = SIGPROC::Pacf(x, minorder, maxorder, betastmp);

	  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) = -betastmp.Rows(1,order(i)).Reverse();
		  betas.SubMatrix(1,order(i),i,i) = betastmp.Rows(1,order(i));
		}
	      
	      Matrix Kinv(sizeTS, sizeTS);
	      Kinv = 0;
  
	      for(int j = 1; j <= sizeTS; j++)
		{
		  if(order(i)==1)
		    {
		      float arone = betastmp(1);
		      for(int k = 1; k <= sizeTS; k++)
			{
			  Kinv(j,k) = pow(arone,abs(k-j));
			}
		    }
		  else
		    Kinv.SubMatrix(j,j,1,j) = Krow.Rows(sizeTS-j+1,sizeTS).t();		  
		}
	     
	      //MISCMATHS::write_ascii_matrix(Kinv,"Kinv");

	      // Kinv now becomes V:
	      if(order(i)!=1)
		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);

      // output betas:
      Log::getInstance().out("betas", betas, false);
      betas.unthresholdSeries(xdata.getDims(),xdata.getPreThresholdPositions());
      betas.writeAsFloat(Log::getInstance().getDir() + "/betas");

      countLargeE = 0;
      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, int usanthresh) {
    Tracer trace("AutoCorrEstimator::spatiallySmooth");
    
    Log& logger = Log::getInstance();

    if(usanthresh == 0)
      {
	// Establish epi thresh to use:
	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 < MISCMATHS::Min(40,int(xdata.getNumVolumes()/4)); 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() + "/usanSize*" << '\0';
    
    cerr << rmfilesstr << endl;

    system(rmfilesstr);

    cerr << "Completed" << endl;
  }

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

    AutoCorr(xdata, acEst, E, 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::multitaper(int M) {
    Tracer tr("AutoCorrEstimator::multitaper");
    
    cerr << "Multitapering... ";
    int sizeTS = xdata.getNumVolumes();
    int numTS = xdata.getNumSeries();

    Matrix slepians;
    getSlepians(M, sizeTS, slepians);

    //Log::getInstance().out("slepians", slepians, false);
    
    ColumnVector x(zeropad);
    x = 0;
    ColumnVector fft_real;
    ColumnVector fft_im;
    ColumnVector dummy(zeropad);
    ColumnVector dummy2;
    ColumnVector realifft(zeropad);
    dummy = 0;
    
    Matrix Sk(zeropad, slepians.Ncols());
    acEst.ReSize(sizeTS, numTS);
    acEst = 0;

    for(int i = 1; i <= numTS; i++) 
      {
	// Compute FFT for each slepian taper
	for(int k = 1; k <= slepians.Ncols(); k++) 
	  {
	    x.Rows(1,sizeTS) = SP(slepians.Column(k), xdata.getSeries(i));
	    
	    //Log::getInstance().out("x", xdata.getSeries(i), false);

	    FFT(x, dummy, fft_real, fft_im);
	    for(int j = 1; j <= zeropad; j++)
	    {
	      // (x+iy)(x-iy) = x^2 + y^2
	      fft_real(j) = fft_real(j)*fft_real(j) + fft_im(j)*fft_im(j);
	      Sk(j,k) = fft_real(j);
	    }
	  }

	// Pool multitaper FFTs
	fft_im = 0;
	for(int j = 1; j <= zeropad; j++)
	  {
	    fft_real(j) = MISCMATHS::mean(ColumnVector(Sk.Row(j).t()));
	    
	  }

	// IFFT to get autocorr
	FFTI(fft_real, fft_im, realifft, dummy2);
	//Log::getInstance().out("Sk", Sk, false);
	//Log::getInstance().out("realifft", realifft);
	//Log::getInstance().out("fftreal", fft_real);
	
	float varx = MISCMATHS::var(ColumnVector(x.Rows(1,sizeTS)));
	acEst.getSeries(i) = realifft.Rows(1,sizeTS)/varx;
      }
    countLargeE = 0;
    cerr << "Completed" << endl;
  }

  void AutoCorrEstimator::getSlepians(int M, int sizeTS, Matrix& slepians) {
    
    Tracer tr("AutoCorrEstimator::getSlepians");
    slepians.ReSize(sizeTS, 2*M);
    
    ifstream in;
    
    char strc[10];
    ostrstream osc(strc,10);
    osc << sizeTS << "_" << M << '\0';
	
    string fname("/usr/people/woolrich/parads/mt_" + string(strc));
    in.open(fname.c_str(), ios::in);
    if(!in)
      throw Exception("Multitapering: Slepians file not found");
   
    for(int j = 1; j <= sizeTS; j++) 
      {
	for(int i = 1; i <= 2*M; i++) 
	  {
	    in >> slepians(j,i);
	  }
      }	

    in.close();
  }

  void AutoCorrEstimator::tukey(int M) {
    
    Tracer tr("AutoCorrEstimator::tukey");

    cerr << "Tukey M = " << M << endl;

    cerr << "Tukey estimates... ";
    int sizeTS = xdata.getNumVolumes();
	
    ColumnVector window(M);

    for(int j = 1; j <= M; j++)
      {
	window(j) = 0.5*(1+cos(M_PI*j/(float(M))));
      }
    
    for(int i = 1; i <= xdata.getNumSeries(); i++) {
	
	acEst.SubMatrix(1,M,i,i) = SP(acEst.SubMatrix(1,M,i,i),window);
	acEst.SubMatrix(M+1,sizeTS,i,i) = 0;
	
    }
    countLargeE = 0;
    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)/float(expwidth),exppow));
	      }	
	  }
	
    }
    countLargeE = 0;

    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(ColumnVector(acEst.getVolume(i).AsColumn()));
	}
    }
}