/*  /*  cluster.cc

    Mark Jenkinson, FMRIB Image Analysis Group

    Copyright (C) 2000 University of Oxford  */

/*  CCOPYRIGHT  */

#include "fmribmain.h"
#include "newimageall.h"
#include <vector>

using namespace NEWIMAGE;

void print_usage(char *argv[])
{
  cerr << "Usage: " << argv[0] << " <input file> <output file> "  
       << "<threshold> <output_format> [-fe]" << endl;
  cerr << "[-f] : threshold is fractional (of robust range) instead of "
       << "absolute"<<endl;
  cerr << "[-e] : used only edge connections (6-way) instead of "
       << "corner connections (26-way)"<<endl;
  cerr << "If threshold is -ve then binarization is inverted"<<endl;
  cerr << "Output formats: the output intensity signifies:"<<endl;
  cerr << "0 - size"<<endl;
  cerr << "1 - size order"<<endl;
  cerr << "2 - highest / lowest original value"<<endl;
  cerr << "3 - mean original value"<<endl;
  //  cerr << "4 - unprocessed unique integer label"<<endl;
}


template <class T>
void get_stats(const volume<int>& labelim, const volume<T>& origim,
	       std::vector<int>& size,
	       std::vector<T>& maxvals, std::vector<T>& meanvals,
	       bool invert) 
{
  int labelnum = labelim.max();
  size.resize(labelnum+1,0);
  maxvals.resize(labelnum+1, (T) 0);
  meanvals.resize(labelnum+1,(T) 0);
  std::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);
	size[idx]++;
	sum[idx]+=(float) origim(x,y,z);
	if (invert) {
	  if ((size[idx]==1) || (maxvals[idx]>origim(x,y,z))) {
	    maxvals[idx] = origim(x,y,z);
	  }
	} else {
	  if ((size[idx]==1) || (maxvals[idx]<origim(x,y,z))) {
	    maxvals[idx] = origim(x,y,z);
	  }
	}
      }
    }
  }
  for (int n=0; n<=labelnum; n++) {
    if (size[n]>0.0)  meanvals[n] = (T) (sum[n]/((float) size[n]));
  }
}

void get_sizeorder(const std::vector<int>& size, std::vector<int>& sizeorder) 
{
  // A BAD O(N^2) SORTING ALGORITHM...
  int labelnum=size.size()-1, maxidx=0, maxsize=0;
  sizeorder.resize(labelnum,0);
  std::vector<int> sizecopy(size);
  for (int n=labelnum; n>=1; n--) {
    maxidx=0;
    maxsize=0;
    for (int m=1; m<=labelnum; m++) {
      if (sizecopy[m]>maxsize) {
	maxidx = m;
	maxsize = sizecopy[m];
      }
    }
    sizecopy[maxidx]=0;
    sizeorder[maxidx] = n;
  }
}


template <class T>
void relabel_image(const volume<int>& labelim, volume<T>& relabelim,
		   const std::vector<T>& 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) = newlabels[labelim(x,y,z)];
      }
    }
  }
}


template <class T>
int fmrib_main(int argc, char *argv[])
{
  volume<int> labelim;
  int outform = atoi(argv[4]);
  std::vector<int> size;
  std::vector<T> maxvals, meanvals;
  {
    // read in the volume
    volume<T> vin;
    read_volume(vin,argv[1]);
    
    // threshold the volume
    bool invert=false;
    float th;
    string opt1="", opt2="", opt="";
    if (argc>=6) { opt = argv[5]; }
    if (opt=="-f") { opt1="-f"; }
    else if (opt=="-e") { opt2="-e"; }
    else if ( (opt=="-fe") || (opt=="-ef") ) { opt1="-f"; opt2="-e"; }
    else if (opt!="") { 
      cerr << "Unrecognised option "<<opt<<endl<<endl;
      print_usage(argv); 
      return -1;
    }

    bool use_corners=true;
    if (opt2=="-e") { use_corners=false; }

    th = atof(argv[3]);
    if (th<0.0) { 
      invert=true;
      th=-th;
    }
    if ( opt1=="-f" ) {
      float frac = th;
      th = frac*(vin.robustmax() - vin.robustmin()) + vin.robustmin();
    }
    vin.threshold((T) th);
    if (invert) { vin = ((T) 1) - vin; }
    
    // Find the connected components
    labelim = connected_components(vin,use_corners);
    
    // process according to the output format
    read_volume(vin,argv[1]);  // re-read the original volume
    get_stats(labelim,vin,size,maxvals,meanvals,invert);
  }

  cout << "CLUSTERS " << size.size()-1 << endl;

  // save the results
  if (outform==0) {
    volume<int> relabeledim;
    size[0] = 0;
    relabel_image(labelim,relabeledim,size);
    save_volume(relabeledim,argv[2]);
  } else if (outform==1) {
    volume<int> relabeledim;
    std::vector<int> sizeorder;
    get_sizeorder(size,sizeorder);
    relabel_image(labelim,relabeledim,sizeorder);
    save_volume(relabeledim,argv[2]);
  } else if (outform==2) {
    volume<T> relabeledim;
    relabel_image(labelim,relabeledim,maxvals);
    save_volume(relabeledim,argv[2]);
  } else if (outform==3) {
    volume<T> relabeledim;
    relabel_image(labelim,relabeledim,meanvals);
    save_volume(relabeledim,argv[2]);
  } else if (outform==4) {
    save_volume(labelim,argv[2]);
  } else {
    cerr << "Unrecognised output format" << endl;
    print_usage(argv);
    return -1;
  }

  return 0;
}



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

  Tracer tr("main");
  if (argc<5) { 
    print_usage(argv);
    return -1; 
  }
  
  return call_fmrib_main(dtype(argv[1]),argc,argv);
}