/*  fsl_glm - 

    Christian F. Beckmann, FMRIB Image Analysis Group

    Copyright (C) 2006-2008 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_glm (Version 1.05)")+
		string("\nCopyright(c) 2004-2008, University of Oxford (Christian F. Beckmann)\n")+
		string(" \n Simple GLM usign ordinary least-squares regression on\n")+
		string(" time courses and/or 3D/4D imges against time courses \n")+
		string(" or 3D/4D images\n\n");
  string examples="fsl_glm -i <input> -d <design> -o <output> [options]";

//Command line Options {
  Option<string> fnin(string("-i,--in"), string(""),
		string("        input file name (text matrix or 3D/4D image file)"),
		true, requires_argument);
  Option<string> fnout(string("-o,--out"), string(""),
		string("output file name for GLM parameter estimates (GLM betas)"),
		true, requires_argument);
  Option<string> fndesign(string("-d,--design"), string(""),
		string("file name of the GLM design matrix (time courses or spatial maps)"),
		true, requires_argument);
  Option<string> fnmask(string("-m,--mask"), string(""),
		string("mask image file name if input is image"),
		false, requires_argument);
  Option<string> fncontrasts(string("-c,--contrasts"), string(""),
		string("matrix of t-statistics contrasts"),
		false, requires_argument);
  Option<string> fnftest(string("-f,--ftests"), string(""),
		string("matrix of F-tests on contrasts"),
		false, requires_argument,false);
	Option<int> dofset(string("--dof"),0,
		string("        set degrees-of-freedom explicitly"),
		false, requires_argument);
	Option<bool> perfvn(string("--vn"),FALSE,
		string("        perfrom MELODIC variance-normalisation on data"),
		false, requires_argument);
	Option<int> help(string("-h,--help"), 0,
		string("display this help text"),
		false,no_argument);
	// Output options	
	Option<string> outcope(string("--out_cope"),string(""),
		string("output file name for COPEs (either as text file or image)"),
		false, requires_argument);
	Option<string> outz(string("--out_z"),string(""),
		string("        output file name for Z-stats (either as text file or image)"),
		false, requires_argument);
	Option<string> outt(string("--out_t"),string(""),
		string("        output file name for t-stats (either as text file or image)"),
		false, requires_argument);
	Option<string> outp(string("--out_p"),string(""),
		string("        output file name for p-values of Z-stats (either as text file or image)"),
		false, requires_argument);
	Option<string> outf(string("--out_f"),string(""),
		string("        output file name for F-value of full model fit"),
		false, requires_argument);
	Option<string> outpf(string("--out_pf"),string(""),
		string("output file name for p-value for full model fit"),
		false, requires_argument);
	Option<string> outres(string("--out_res"),string(""),
		string("output file name for residuals"),
		false, requires_argument);
	Option<string> outvarcb(string("--out_varcb"),string(""),
		string("output file name for variance of COPEs"),
		false, requires_argument);
	Option<string> outsigsq(string("--out_sigsq"),string(""),
		string("output file name for residual noise variance sigma-square"),
		false, requires_argument);
	Option<string> outdata(string("--out_data"),string(""),
		string("output file name for pre-processed data"),
		false, requires_argument);
	Option<string> outvnscales(string("--out_vnscales"),string(""),
		string("output file name for scaling factors for variance normalisation"),
		false, requires_argument);
		/*
}
*/
//Globals {
	Melodic::basicGLM glm;
	int voxels = 0;
	Matrix data;
	Matrix design;
	Matrix contrasts;
	Matrix fcontrasts;
	Matrix meanR;
	RowVector vnscales;
	volume<float> mask;
	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 if(fsl_imageexists(fndesign.value()))
		write_ascii_matrix(what.t(),fname);
	else
		write_ascii_matrix(what,fname);
}

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{
			mask = tmpdata[0]*0.0+1.0;	
		}
		
		data = tmpdata.matrix(mask);
		voxels = data.Ncols();
	}
	else
		data = read_ascii_matrix(fnin.value());	

	if(fsl_imageexists(fndesign.value())){//read design
		volume4D<float> tmpdata;
		read_volume4D(tmpdata,fndesign.value());
		if(!samesize(tmpdata[0],mask)){
			cerr << "ERROR: GLM design does not match input image in size" << endl;
			return 1;
		}
		design = tmpdata.matrix(mask).t();
		data = data.t();
	}else{
		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);
	if(fncontrasts.value()>""){//read contrast		
		contrasts = read_ascii_matrix(fncontrasts.value());
		if(!(contrasts.Ncols()==design.Ncols())){
			cerr << "ERROR: contrast matrix GLM design does not match GLM design" << endl;
			return 1;
		}
	}else{
		contrasts = IdentityMatrix(design.Ncols());
		contrasts &= -1.0 * contrasts;
	}
	return 0;	
}

void write_res(){	
	if(fnout.value()>"")
		saveit(glm.get_beta(),fnout.value());
	if(outcope.value()>"")
		saveit(glm.get_cbeta(),outcope.value());
	if(outz.value()>"")
		saveit(glm.get_z(),outz.value());
	if(outt.value()>"")
		saveit(glm.get_t(),outt.value());
	if(outp.value()>"")
		saveit(glm.get_p(),outp.value());
	if(outf.value()>"")
		saveit(glm.get_f_fmf(),outf.value());
	if(outpf.value()>"")
		saveit(glm.get_pf_fmf(),outpf.value());
	if(outres.value()>"")
		saveit(glm.get_residu(),outres.value());
	if(outvarcb.value()>"")
		saveit(glm.get_varcb(),outvarcb.value());
	if(outsigsq.value()>"")
		saveit(glm.get_sigsq(),outsigsq.value());
	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);

	glm.olsfit(data,design,contrasts,dofset.value());
	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(fncontrasts);
			options.add(fnmask);
			options.add(fnftest);
			options.add(dofset);
			options.add(perfvn);
			options.add(help);
			options.add(outcope);
			options.add(outz);
			options.add(outt);
			options.add(outp);
			options.add(outf);
			options.add(outpf);
			options.add(outres);
			options.add(outvarcb);
			options.add(outsigsq);
			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;
	  } 
	}