/*  fsl_regfilt - 

    Christian Beckmann, FMRIB Image Analysis Group

    Copyright (C) 2006-2007 University of Oxford  */

/*  CCOPYRIGHT  */

#include "libvis/miscplot.h"
#include "miscmaths/miscmaths.h"
#include "miscmaths/miscprob.h"
#include "utils/options.h"
#include <vector>
#include "newimage/newimageall.h"
#include "melhlprfns.h"

using namespace MISCPLOT;
using namespace MISCMATHS;
using namespace Utilities;
using namespace std;

// The two strings below specify the title and example usage that is
// printed out as the help or usage message

  string title=string("fsl_regfilt (Version 1.0)")+
		string("\n\n Copyright(c) 2007, University of Oxford (Christian F. Beckmann)\n")+
		string(" Data de-noising by regressing out part of a design matrix\n")+
		string(" using simple OLS regression on 4D images");
  string examples="fsl_regfilt -i <input> -d <design> -f <components> -o <out> [options]";

//Command line Options {
  Option<string> fnin(string("-i,--in"), string(""),
		string("        input file name (4D image)"),
		true, requires_argument);
  Option<string> fnout(string("-o,--out"), string(""),
		string("output file name for the filtered data"),
		true, requires_argument);
  Option<string> fndesign(string("-d,--design"), string(""),
		string("file name of the GLM design matrix (time courses)"),
		true, requires_argument);
  Option<string> fnmask(string("-m,--mask"), string(""),
		string("mask image file name"),
		false, requires_argument);
	Option<string> filter(string("-f,--filter"),string(""),
		string("filter out part of the regression model, e.g. -f \"1,2,3\""),
		true, requires_argument);
	Option<bool> verbose(string("-v"),FALSE,
		string("        switch on diagnostic messages"),
		false, no_argument);
	Option<bool> perfvn(string("--vn"),FALSE,
		string("        perfrom variance-normalisation on data"),
		false, no_argument);
	Option<int> help(string("-h,--help"), 0,
		string("display this help text"),
		false,no_argument);
	// Output options	
	Option<string> outdata(string("--out_data"),string(""),
		string("output data"),
		false, requires_argument);
	Option<string> outvnscales(string("--out_vnscales"),string(""),
		string("output scaling factors for variance normalisation"),
		false, requires_argument);
		/*
}
*/
//Globals {
	int voxels = 0;
	Matrix data;
	Matrix design;
	Matrix meanR;
	RowVector vnscales;
	volume<float> mask;
	volume<float> Mean;
	volumeinfo volinf;  /*
}
*/
////////////////////////////////////////////////////////////////////////////

// Local functions
void save4D(Matrix what, string fname){
		if(what.Ncols()==data.Ncols()||what.Nrows()==data.Nrows()){
			volume4D<float> tempVol;
			if(what.Nrows()>what.Ncols())
				tempVol.setmatrix(what.t(),mask);
			else
				tempVol.setmatrix(what,mask);
			save_volume4D(tempVol,fname,volinf);
		}
}

bool isimage(Matrix what){
	if((voxels > 0)&&(what.Ncols()==voxels || what.Nrows()==voxels))
		return TRUE;
	else
		return FALSE;
}

void saveit(Matrix what, string fname){
	if(isimage(what))
		save4D(what,fname);
	else
		write_ascii_matrix(what,fname);
}

int dofilter(){
	if(!isimage(data)){
		cerr << "ERROR: need to specify 4D input to use filtering" << endl;
		return 1;
	}
	if(verbose.value())
		cout << "  Calculating maps " << endl;  
	Matrix unmixMatrix = pinv(design);
  Matrix maps = unmixMatrix * data;

  Matrix noisedes;
  Matrix noisemaps;

  int ctr=0;    
  char *p;
  char t[1024];
  const char *discard = ", [];{(})abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ~!@#$%^&*_-=+|\':><./?";
  
  strcpy(t, filter.value().c_str());
  p=strtok(t,discard);
  ctr = atoi(p);
  if(ctr>0 && ctr<=design.Ncols()){
    noisedes = design.Column(ctr);
    noisemaps  = maps.Row(ctr).t();    
  }

  do{
    p=strtok(NULL,discard);
    if(p){
			ctr = atoi(p);
      if(ctr>0 && ctr<=design.Ncols()){
  			noisedes |= design.Column(ctr);
  			noisemaps  |= maps.Row(ctr).t();
			}
    }
  }while(p);
  Matrix newData;
	if(verbose.value())
		cout << "  Calculating filtered data " << endl;
 	newData = data - noisedes * noisemaps.t();
  newData = newData + ones(newData.Nrows(),1)*meanR;
  
	save4D(newData,fnout.value());
  return 0;	
}

int setup(){
	if(fsl_imageexists(fnin.value())){//read data
		//input is 3D/4D vol
		volume4D<float> tmpdata;
		read_volume4D(tmpdata,fnin.value(),volinf);
		
		// create mask
		if(fnmask.value()>""){
			read_volume(mask,fnmask.value());
			if(!samesize(tmpdata[0],mask)){
				cerr << "ERROR: Mask image does not match input image" << endl;
				return 1;
			};
		}else{
			if(verbose.value())
				cout << "  Creating mask image "  << endl;
			Mean = meanvol(tmpdata);
			float Mmin, Mmax;
			Mmin = Mean.min(); Mmax = Mean.max();
			mask = binarise(Mean,float(Mmin + 0.01* (Mmax-Mmin)),Mmax);
		}
		
		data = tmpdata.matrix(mask);
		voxels = data.Ncols();
		if(verbose.value())
			cout << " Data matrix size : " << data.Nrows() << " x " << voxels << endl;
		
	}else{
		cerr << "ERROR: cannot read input image " << fnin.value()<<endl;
		return 1;
	}

	design = read_ascii_matrix(fndesign.value());

	meanR=mean(data,1);
	data = remmean(data,1);
	design = remmean(design,1);
	if(perfvn.value())
		vnscales = Melodic::varnorm(data);
	return 0;	
}

void write_res(){	
	if(outdata.value()>"")
		saveit(data,outdata.value());
	if(outvnscales.value()>"")
		saveit(vnscales,outvnscales.value());
}

int do_work(int argc, char* argv[]) {
  if(setup())
		exit(1);

	if(dofilter())
		exit(1);	
	write_res();
	return 0;
}

////////////////////////////////////////////////////////////////////////////

int main(int argc,char *argv[]){
	  Tracer tr("main");
	  OptionParser options(title, examples);
	  try{
	    // must include all wanted options here (the order determines how
	    //  the help message is printed)
			options.add(fnin);
			options.add(fnout);
			options.add(fndesign);
			options.add(fnmask);
			options.add(filter);
			options.add(perfvn);
			options.add(verbose);
			options.add(help);
			options.add(outdata);
			options.add(outvnscales);
	    options.parse_command_line(argc, argv);

	    // line below stops the program if the help was requested or 
	    //  a compulsory option was not set
	    if ( (help.value()) || (!options.check_compulsory_arguments(true)) ){
				options.usage();
				exit(EXIT_FAILURE);
	    }else{
	  		// Call the local functions
	  		return do_work(argc,argv);
			}
		}catch(X_OptionError& e) {
			options.usage();
	  	cerr << endl << e.what() << endl;
	    exit(EXIT_FAILURE);
	  }catch(std::exception &e) {
	    cerr << e.what() << endl;
	  } 
}