/*  cluster.cc

    Mark Jenkinson, FMRIB Image Analysis Group

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

/*  CCOPYRIGHT */

#include <vector>
#include <algorithm>
#include <iomanip>
#include "newimage/fmribmain.h"
#include "newimage/newimageall.h"
#include "utils/options.h"
#include "infer.h"

#define _GNU_SOURCE 1
#define POSIX_SOURCE 1

using namespace NEWIMAGE;
using std::vector;
using namespace Utilities;

string title="cluster (Version 1.2)\nCopyright(c) 2000-2002, University of Oxford (Mark Jenkinson)";
string examples="cluster --in=<filename> --thresh=<value> [options]";

Option<bool> verbose(string("-v,--verbose"), false, 
		     string("switch on diagnostic messages"), 
		     false, no_argument);
Option<bool> help(string("-h,--help"), false,
		  string("display this message"),
		  false, no_argument);
Option<bool> medxcoords(string("-m,--medx"), false,
		  string("use MEDx coordinate conventions"),
		  false, no_argument);
Option<bool> mm(string("--mm"), false,
		  string("use mm, not voxel, coordinates"),
		  false, no_argument);
Option<bool> minv(string("--min"), false,
		  string("find minima instead of maxima"),
		  false, no_argument);
Option<bool> fractional(string("--fractional"), false,
			string("interprets the threshold as a fraction of the robust range"),
			false, no_argument);
Option<bool> no_table(string("--no_table"), false,
		      string("suppresses printing of the table info"),
		      false, no_argument);
Option<bool> minclustersize(string("--minclustersize"), false,
		      string("prints out minimum significant cluster size"),
		      false, no_argument);
Option<int> voxvol(string("--volume"), 0,
		   string("number of voxels in the mask"),
		   false, requires_argument);
Option<int> numconnected(string("--connectivity"), 26,
		   string("the connectivity of voxels (default 26)"),
		   false, requires_argument);
Option<int> mx_cnt(string("-n,--num"), 6,
		   string("no of local maxima to report"),
		   false, requires_argument);
Option<float> dLh(string("-d,--dlh"), 4.0,
		  string("smoothness estimate = sqrt(det(Lambda))"),
		  false, requires_argument);
Option<float> thresh(string("-t,--thresh,--zthresh"), 2.3,
		     string("threshold for input volume"),
		     true, requires_argument);
Option<float> pthresh(string("-p,--pthresh"), 0.01,
		      string("p-threshold for clusters"),
		      false, requires_argument);
Option<string> inputname(string("-i,--in,-z,--zstat"), string(""),
			 string("filename of input volume"),
			 true, requires_argument);
Option<string> copename(string("-c,--cope"), string(""),
			string("filename of input cope volume"),
			false, requires_argument);
Option<string> outindex(string("-o,--oindex"), string(""),
			string("filename for output of cluster index (in size order)"),
			false, requires_argument);
Option<string> outlmax(string("--olmax"), string(""),
			string("filename for output of local maxima text file"),
			false, requires_argument);
Option<string> outthresh(string("--othresh"), string(""),
			 string("filename for output of thresholded image"),
			 false, requires_argument);
Option<string> outsize(string("--osize"), string(""),
		       string("filename for output of size image"),
		       false, requires_argument);
Option<string> outmax(string("--omax"), string(""),
		      string("filename for output of max image"),
		      false, requires_argument);
Option<string> outmean(string("--omean"), string(""),
		       string("filename for output of mean image"),
		       false, requires_argument);
Option<string> transformname(string("-x,--xfm"), string(""),
		       string("filename for transform of input->talairach volume"),
		       false, requires_argument);
Option<string> talvolname(string("--talvol"), string(""),
		       string("filename for talairach volume"),
		       false, requires_argument);



int num(const char x) { return (int) x; }
short int num(const short int x) { return x; }
int num(const int x) { return x; }
float num(const float x) { return x; }
double num(const double x) { return x; }


template <class T>
struct triple { T x; T y; T z; };

template <class T, class S>
void copyconvert(const vector<triple<T> >& oldcoords, 
		 vector<triple<S> >& newcoords)
{
  newcoords.erase(newcoords.begin(),newcoords.end());
  newcoords.resize(oldcoords.size());
  for (unsigned int n=0; n<oldcoords.size(); n++) {
    newcoords[n].x = (S) oldcoords[n].x;
    newcoords[n].y = (S) oldcoords[n].y;
    newcoords[n].z = (S) oldcoords[n].z;
  }
}

template <class T, class S>
void vox2medxcoord(vector<triple<T> >& coords, const volume<S>& vol)
{
  int ymax = vol.ysize();
  for (unsigned int n=0; n<coords.size(); n++) {
    coords[n].y = (T) (ymax - 1) - coords[n].y;
  }
}

void vox2mmcoord(vector<triple<float> >& coords, 
				   float vx, float vy, float vz)
{
  for (unsigned int n=0; n<coords.size(); n++) {
    coords[n].x = coords[n].x * vx;
    coords[n].y = coords[n].y * vy;
    coords[n].z = coords[n].z * vz;
  }
}


template <class S>
void vox2mmcoord(vector<triple<float> >& coords, 
				   const volume<S>& vol)
{
  float vx = vol.xdim()*vol.xsign(), vy = vol.ydim()*vol.ysign(), 
    vz = vol.zdim()*vol.zsign();
  vox2mmcoord(coords,vx,vy,vz);
}


template <class T, class S>
void addoriginoffset(vector<triple<T> >& coords, const volume<S>& vol)
{
  ColumnVector orig;
  orig = vol.getorigin();
  if (orig.SumAbsoluteValue() < 0.1)  return;  // zero origin -> no work
  for (unsigned int n=0; n<coords.size(); n++) {
    coords[n].x -= (T) (orig(1) - 1.0);   // -1.0 convert from SPM convention
    coords[n].y -= (T) (orig(2) - 2.0);   // -2.0 convert from SPM convention (NFI)
    coords[n].z -= (T) (orig(3) - 1.0);   // -1.0 convert from SPM convention
  }
}


template <class T, class S>
void applyvoxelxfm(vector<triple<T> >& coords, const Matrix& mat, 
	      const volume<S>& source, const volume<S>& dest)
{
  Matrix voxmat;
  voxmat = dest.sampling_mat().i() * mat * source.sampling_mat();
  ColumnVector vec(4);
  for (unsigned int n=0; n<coords.size(); n++) {
    vec << coords[n].x << coords[n].y << coords[n].z << 1.0;
    vec = voxmat * vec;     // apply voxel xfm
    coords[n].x = vec(1);
    coords[n].y = vec(2);
    coords[n].z = vec(3);
  }
}


template <class T>
void get_stats(const volume<int>& labelim, const volume<T>& origim,
	       vector<int>& size,
	       vector<T>& maxvals, vector<float>& meanvals,
	       vector<triple<int> >& max, vector<triple<float> >& cog,
	       bool minv) 
{
  int labelnum = labelim.max();
  size.resize(labelnum+1,0);
  maxvals.resize(labelnum+1, (T) 0);
  meanvals.resize(labelnum+1,0.0f);
  triple<int> zero;
  zero.x = 0; zero.y = 0; zero.z = 0;
  triple<float> zerof;
  zerof.x = 0; zerof.y = 0; zerof.z = 0;
  max.resize(labelnum+1,zero);
  cog.resize(labelnum+1,zerof);
  vector<float> sum(labelnum+1,0.0);
  for (int z=labelim.minz(); z<=labelim.maxz(); z++) {
    for (int y=labelim.miny(); y<=labelim.maxy(); y++) {
      for (int x=labelim.minx(); x<=labelim.maxx(); x++) {
	int idx = labelim(x,y,z);
	T oxyz = origim(x,y,z);
	size[idx]++;
	cog[idx].x+=((float) oxyz)*x;
	cog[idx].y+=((float) oxyz)*y;
	cog[idx].z+=((float) oxyz)*z;
	sum[idx]+=(float) oxyz;
	if ( (size[idx]==1) || 
	     ( (oxyz>maxvals[idx]) && (!minv) ) || 
	     ( (oxyz<maxvals[idx]) && (minv) ) ) 
	  {
	    maxvals[idx] = oxyz;
	    max[idx].x = x;
	    max[idx].y = y;
	    max[idx].z = z;
	  }
      }
    }
  }
  for (int n=0; n<=labelnum; n++) {
    if (size[n]>0.0) {
      meanvals[n] = (sum[n]/((float) size[n]));
    }
    if (sum[n]>0.0) {
      cog[n].x /= sum[n];
      cog[n].y /= sum[n];
      cog[n].z /= sum[n];
    }
  }
}


vector<int> get_sortindex(const vector<int>& vals)
{
  // return the mapping of old indices to new indices in the
  //   new *ascending* sort of vals
  int length=vals.size();
  vector<pair<int, int> > sortlist(length);
  for (int n=0; n<length; n++) {
    sortlist[n] = pair<int, int>(vals[n],n);
  }
  sort(sortlist.begin(),sortlist.end());  // O(N.log(N))
  vector<int> idx(length);
  for (int n=0; n<length; n++) {
    idx[n] = sortlist[n].second;
  }
  return idx;
}


void get_sizeorder(const vector<int>& size, vector<int>& sizeorder) 
{
  vector<int> sizecopy(size), idx;
  idx = get_sortindex(sizecopy);

  // second part of pair is now the prior-index of the sorted values
  int length = size.size();
  sizeorder.resize(length,0);
  for (int n=0; n<length; n++) {
    sizeorder[idx[n]] = n;    // maps old index to new
  }
}


template <class T, class S>
void relabel_image(const volume<int>& labelim, volume<T>& relabelim,
		   const vector<S>& newlabels)
{
  copyconvert(labelim,relabelim);
  for (int z=relabelim.minz(); z<=relabelim.maxz(); z++) {
    for (int y=relabelim.miny(); y<=relabelim.maxy(); y++) {
      for (int x=relabelim.minx(); x<=relabelim.maxx(); x++) {
	relabelim(x,y,z) = (T) newlabels[labelim(x,y,z)];
      }
    }
  }
}


template <class T, class S>
void print_results(const vector<int>& idx, 
		   const vector<int>& size,
		   const vector<int>& pthreshindex,
		   const vector<float>& pvals,
		   const vector<float>& logpvals,
		   const vector<T>& maxvals,
		   const vector<triple<S> >& maxpos,
		   const vector<triple<float> >& cog,
		   const vector<T>& copemaxval,
		   const vector<triple<S> >& copemaxpos,
		   const vector<float>& copemean,
		   const volume<T>& zvol, const volume<T>& cope, 
		   const volume<int> &labelim)
{
  int length=size.size();

  vector<triple<float> > fmaxpos, fcopemaxpos, fcog;
  copyconvert(maxpos,fmaxpos);
  copyconvert(cog,fcog);
  copyconvert(copemaxpos,fcopemaxpos);
  volume<T> talvol;
  Matrix trans;
  const volume<T> *refvol = &zvol;
  if ( !transformname.unset() && !talvolname.unset() ) {
    read_volume(talvol,talvolname.value());
    read_matrix(trans,transformname.value(),zvol);
    if (verbose.value()) { 
      cout << "Transformation Matrix filename = "<<transformname.value()<<endl;
      cout << trans.Nrows() << " " << trans.Ncols() << endl; 
      cout << "Transformation Matrix = " << endl;
      cout << trans << endl;
    }
    applyvoxelxfm(fmaxpos,trans,zvol,talvol);
    applyvoxelxfm(fcog,trans,zvol,talvol);
    if (!copename.unset()) applyvoxelxfm(fcopemaxpos,trans,zvol,talvol);
  }
  if ( medxcoords.value() && (!mm.value()) ) {
    vox2medxcoord(fmaxpos,zvol);
    vox2medxcoord(fcog,zvol);
    if (!copename.unset()) vox2medxcoord(fcopemaxpos,cope);
  }
  if (mm.value()) {
    if ( !transformname.unset() && !talvolname.unset() ) refvol = &talvol;
    if (verbose.value()) { 
      cout << "Using origin : " << (refvol->getorigin()).t() << endl; 
    }
    addoriginoffset(fmaxpos,*refvol);
    addoriginoffset(fcog,*refvol);
    if (!copename.unset()) addoriginoffset(fcopemaxpos,*refvol);
    vox2mmcoord(fmaxpos,*refvol);
    vox2mmcoord(fcog,*refvol);
    if (!copename.unset()) vox2mmcoord(fcopemaxpos,*refvol);  // used cope before
  }

  string units=" (vox)";
  if (!mm.unset()) units=" (mm)";
  string tablehead;
  tablehead = "Cluster Index\tVoxels";
  if (!pthresh.unset()) tablehead += "\tP\t-log10(P)";
  tablehead += "\tMax Z\tx" + units + "\ty" + units + "\tz" + units
    + "\tCOG x" + units + "\tCOG y" + units + "\tCOG z" + units;
  if (!copename.unset()) {
    tablehead+= "\tMax COPE\tx" + units + "\ty" + units + "\tz" + units 
                 + "\tMean COPE";
  }

  cout << tablehead << endl;

  for (int n=length-1; n>=1; n--) {
    int index=idx[n];
    int k = size[index];
    float p = pvals[index];
    if (pthreshindex[index]>0) {
      float mlog10p;
      mlog10p = -logpvals[index];
      cout << setprecision(3) << num(pthreshindex[index]) << "\t" << k << "\t"; 
      if (!pthresh.unset()) { cout << num(p) << "\t" << num(mlog10p) << "\t"; }
      cout << num(maxvals[index]) << "\t" 
	   << num(fmaxpos[index].x) << "\t" << num(fmaxpos[index].y) << "\t" 
	   << num(fmaxpos[index].z) << "\t"
	   << num(fcog[index].x) << "\t" << num(fcog[index].y) << "\t" 
	   << num(fcog[index].z);
      if (!copename.unset()) { 
	cout   << "\t" << num(copemaxval[index]) << "\t"
	       << num(fcopemaxpos[index].x) << "\t" << num(fcopemaxpos[index].y) << "\t" 
	       << num(fcopemaxpos[index].z) << "\t" << num(copemean[index]);
	}
      cout << endl;
    }
  }

  // output local maxima
  if (!outlmax.unset()) {
    ofstream lmaxfile(outlmax.value().c_str());
    if (!lmaxfile)
      cerr << "Could not open file " << outlmax.value() << " for writing" << endl;
    lmaxfile << "Cluster Index\tZ\tx\ty\tz\t" << endl;
    zvol.setextrapolationmethod(zeropad);
    for (int n=size.size()-1; n>=1; n--) {
      int index=idx[n];
      if (pthreshindex[index]>0) {
	vector<int>   lmaxlistZ(size[index]);
	vector<triple<float> > lmaxlistR(size[index]);
	int lmaxlistcounter=0;
	for (int z=labelim.minz(); z<=labelim.maxz(); z++)
	  for (int y=labelim.miny(); y<=labelim.maxy(); y++)
	    for (int x=labelim.minx(); x<=labelim.maxx(); x++)
	      if ( index==labelim(x,y,z) &&
		   zvol(x,y,z)>zvol(x-1,y-1,z-1) &&
		   zvol(x,y,z)>zvol(x,  y-1,z-1) &&
		   zvol(x,y,z)>zvol(x+1,y-1,z-1) &&
		   zvol(x,y,z)>zvol(x-1,y,  z-1) &&
		   zvol(x,y,z)>zvol(x,  y,  z-1) &&
		   zvol(x,y,z)>zvol(x+1,y,  z-1) &&
		   zvol(x,y,z)>zvol(x-1,y+1,z-1) &&
		   zvol(x,y,z)>zvol(x,  y+1,z-1) &&
		   zvol(x,y,z)>zvol(x+1,y+1,z-1) &&
		   zvol(x,y,z)>zvol(x-1,y-1,z) &&
		   zvol(x,y,z)>zvol(x,  y-1,z) &&
		   zvol(x,y,z)>zvol(x+1,y-1,z) &&
		   zvol(x,y,z)>zvol(x-1,y,  z) &&
		   zvol(x,y,z)>=zvol(x+1,y,  z) &&
		   zvol(x,y,z)>=zvol(x-1,y+1,z) &&
		   zvol(x,y,z)>=zvol(x,  y+1,z) &&
		   zvol(x,y,z)>=zvol(x+1,y+1,z) &&
		   zvol(x,y,z)>=zvol(x-1,y-1,z+1) &&
		   zvol(x,y,z)>=zvol(x,  y-1,z+1) &&
		   zvol(x,y,z)>=zvol(x+1,y-1,z+1) &&
		   zvol(x,y,z)>=zvol(x-1,y,  z+1) &&
		   zvol(x,y,z)>=zvol(x,  y,  z+1) &&
		   zvol(x,y,z)>=zvol(x+1,y,  z+1) &&
		   zvol(x,y,z)>=zvol(x-1,y+1,z+1) &&
		   zvol(x,y,z)>=zvol(x,  y+1,z+1) &&
		   zvol(x,y,z)>=zvol(x+1,y+1,z+1) ) {
		lmaxlistZ[lmaxlistcounter]=(int)(1000.0*zvol(x,y,z));
		lmaxlistR[lmaxlistcounter].x=x;
		lmaxlistR[lmaxlistcounter].y=y;
		lmaxlistR[lmaxlistcounter].z=z;
		lmaxlistcounter++;
	      }

	lmaxlistZ.resize(lmaxlistcounter);
	vector<int> lmaxidx = get_sortindex(lmaxlistZ);

	if ( !transformname.unset() && !talvolname.unset() )
	  applyvoxelxfm(lmaxlistR,trans,zvol,talvol);
	if ( medxcoords.value() && (!mm.value()) )
	  vox2medxcoord(lmaxlistR,zvol);
	if (mm.value()) {
	  addoriginoffset(lmaxlistR,*refvol);
	  vox2mmcoord(lmaxlistR,*refvol);
	}

	for (int i=lmaxlistcounter-1; i>MISCMATHS::Max(lmaxlistcounter-1-mx_cnt.value(),-1); i--)
	  lmaxfile << setprecision(3) << pthreshindex[index] << "\t" << lmaxlistZ[lmaxidx[i]]/1000.0 << "\t" << 
	    lmaxlistR[lmaxidx[i]].x << "\t" << lmaxlistR[lmaxidx[i]].y << "\t" << lmaxlistR[lmaxidx[i]].z << endl;
      }
    }
    lmaxfile.close();
  }

}



template <class T>
int fmrib_main(int argc, char *argv[])
{
  volume<int> labelim;
  vector<int> size;
  vector<triple<int> > maxpos;
  vector<triple<float> > cog;
  vector<T> maxvals;
  vector<float> meanvals;
  float th = thresh.value();

  // read in the volume
  volume<T> zvol, mask, cope;
  read_volume(zvol,inputname.value());
  if (verbose.value())  print_volume_info(zvol,"Zvol");
  
  if ( fractional.value() ) {
    float frac = th;
    th = frac*(zvol.robustmax() - zvol.robustmin()) + zvol.robustmin();
  }
  mask = zvol;
  mask.binarise((T) th);
  if (minv.value()) { mask = ((T) 1) - mask; }
  if (verbose.value())  print_volume_info(mask,"Mask");
  
  // Find the connected components
  labelim = connected_components(mask,numconnected.value());
  if (verbose.value())  print_volume_info(labelim,"Labelim");
  
  // process according to the output format
  get_stats(labelim,zvol,size,maxvals,meanvals,maxpos,cog,minv.value());
  if (verbose.value()) {cout<<"Number of labels = "<<size.size()<<endl;}

  // process the cope image if entered
  vector<int> copesize;
  vector<triple<int> > copemaxpos;
  vector<triple<float> > copecog;
  vector<T> copemaxval;
  vector<float> copemean;
  if (!copename.unset()) {
    read_volume(cope,copename.value());
    get_stats(labelim,cope,copesize,copemaxval,copemean,copemaxpos,copecog,
	      minv.value());
  }

  // re-threshold for p
  int length = size.size();
  size[0] = 0;  // force background to be ordered last
  vector<int> pthreshsize;
  vector<float> pvals(length), logpvals(length);
  pthreshsize = size;
  int nozeroclust=0;
  if (!pthresh.unset()) {
    if (verbose.value()) {cout<<"Re-thresholding with p-value"<<endl;}
    Infer infer(dLh.value(), th, voxvol.value());
    if (labelim.zsize()<=1) { infer.setD(2); } // the 2D option
    if (minclustersize.value()) {
      float pmin=1.0;
      int nmin=1;
      while (pmin>=pthresh.value()) {
	pmin=exp(infer(nmin++));
      }
      nmin--;  // now nmin has the value of the smallest cluster under thresh
      cout << "Minimum cluster size under p-threshold = " << nmin << endl;
    }
    for (int n=1; n<length; n++) {
      int k = size[n];
      logpvals[n] = infer(k)/log(10);
      pvals[n] = exp(logpvals[n]*log(10));
      if (pvals[n]>pthresh.value()) {
	pthreshsize[n] = 0;
	nozeroclust++;
      }
    }
  }
  if (verbose.value()) {cout<<"Number of sub-p clusters = "<<nozeroclust<<endl;}

  // get sorted index (will revert to cluster size if no pthresh)
  vector<int> idx;
  idx = get_sortindex(pthreshsize);
  if (verbose.value()) {cout<<pthreshsize.size()<<" labels in sortedidx"<<endl;}

  vector<int> threshidx(length);
  for (int n=length-1; n>=1; n--) {
    int index=idx[n];
    if (pthreshsize[index]>0) {
      threshidx[index] = n - nozeroclust;
    }
  }

  // print table
  if (!no_table.value()) {
    print_results(idx, size, threshidx, pvals, logpvals, maxvals, maxpos, cog,
		  copemaxval, copemaxpos, copemean, zvol, cope, labelim);
  }

  // save relevant volumes
  if (!outindex.unset()) {
    volume<int> relabeledim;
    relabel_image(labelim,relabeledim,threshidx);
    save_volume(relabeledim,outindex.value());
  }
  if (!outsize.unset()) {
    volume<int> relabeledim;
    relabel_image(labelim,relabeledim,pthreshsize);
    save_volume(relabeledim,outsize.value());
  }
  if (!outmax.unset()) {
    volume<T> relabeledim;
    relabel_image(labelim,relabeledim,maxvals);
    save_volume(relabeledim,outmax.value());
  }
  if (!outmean.unset()) {
    volume<float> relabeledim;
    relabel_image(labelim,relabeledim,meanvals);
    save_volume(relabeledim,outmean.value());
  }
  if (!outthresh.unset()) {
    // Threshold the input volume st it is 0 for all non-clusters
    //   and maintains the same values otherwise
    volume<T> lcopy;
    relabel_image(labelim,lcopy,threshidx);
    lcopy.binarise(1);
    save_volume(zvol * lcopy,outthresh.value());
  }

  return 0;
}



int main(int argc,char *argv[])
{

  Tracer tr("main");

  OptionParser options(title, examples);

  try {
    options.add(inputname);
    options.add(thresh);
    options.add(outindex);
    options.add(outthresh);
    options.add(outlmax);
    options.add(outsize);
    options.add(outmax);
    options.add(outmean);
    options.add(pthresh);
    options.add(copename);
    options.add(voxvol);
    options.add(dLh);
    options.add(fractional);
    options.add(numconnected);
    options.add(medxcoords);
    options.add(mm);
    options.add(minv);
    options.add(no_table);
    options.add(minclustersize);
    options.add(transformname);
    options.add(talvolname);
    options.add(mx_cnt);
    options.add(verbose);
    options.add(help);
    
    options.parse_command_line(argc, argv);

    if ( (help.value()) || (!options.check_compulsory_arguments(true)) )
      {
	options.usage();
	exit(EXIT_FAILURE);
      }
    
    if ( (!pthresh.unset()) && (dLh.unset() || voxvol.unset()) ) 
      {
	options.usage();
	cerr << endl 
	     << "Both --dlh and --volume MUST be set if --pthresh is used." 
	     << endl;
	exit(EXIT_FAILURE);
      }
    
    if ( ( !transformname.unset() && talvolname.unset() ) ||
	 ( transformname.unset() && (!talvolname.unset()) ) ) 
      {
	options.usage();
	cerr << endl 
	     << "Both --xfm and --talvol MUST be set together." 
	     << endl;
	exit(EXIT_FAILURE);
      }
  }  catch(X_OptionError& e) {
    options.usage();
    cerr << endl << e.what() << endl;
    exit(EXIT_FAILURE);
  } catch(std::exception &e) {
    cerr << e.what() << endl;
  } 

  return call_fmrib_main(dtype(inputname.value()),argc,argv);

}