#!/usr/local/bin/perl

use strict;
use warnings;
use Cwd;

unless ($ARGV[0]){
    print "sm_analysis.pl version 1.0 (5/3/07) - szymon\@nmr.mgh.harvard.edu\n";	
    print "sm_analysis.pl version 1.1 (5/25/07) - modifications by levy\@nmr.mgh.harvard.edu\n";
    print "\n";
    print "Wrapper for MEG analysis utilizing tools of the MNE stream and some of my own scripts. Takes raw data as input and produces forward and inverse solutions.\n";
    print "\n";
    print "Usage: sm_analysis [options]\n";
    print "\n";
    print "Options:\n";
    print "--runs		:	list numbers of runs to analyze (defaults to 01 02 03 04 05 06 012 011)\n";
    print "--conds		:	list of names of conditions\n";
    print "--old		:	list of trigger numbers to be replaces in the lst file\n";
    print "--new		:	list of trigger numbers to that replace the above\n";
    print "--trigs		:	list of triggers corresponding to conditions\n";
    print "--sets		:	list of data sets in fif files corresponding to conditions\n";
    print "--analysis	:	name of analysis\n";
    print "--type		:	type of analysis, suffix in lst file (e.g. stim, resp, end, fix)\n";
    print "--list		:	path to subject list\n";
    print "--over		:	whether to clear existing analysis directories\n";
    print "--gradonly	:	excludes magnetometers from analysis\n";
    print "--baselines 	:	whether to calculate baselines\n";
    print "--tmin		:	beginning timepoint in ms\n";
    print "--tmax		:	end timepoint in ms\n";
    print "--btmin		:	beginning timepoint of baseline in ms (defaults to -200)\n";
    print "--btmax		:	end timepoint of baseline in ms (defaults to 0)\n";
    print "--nocov 		:	do not calculate covariance matrices, forward solutions and inverse operator\n";
    print "--parent		:	name of parent anlysis; this is where baselines and inverse operator may be copied from\n";
    print "--bad_file	:	path to bad channels file from \$MEG_DIR\n";
    print "--sss	        :	use SSS (Maxwell filtering) instead of SSP\n";
    print "--st_file        :       (only available for sss) read file indicating which subjects to run temporal SSS on.";
    print "--skip_file      :	(only available for sss) read skip file to exclude sections temporally. \n";	
    print "\n";
    exit;
}

# set defaults here:

my $MEG_DIR=$ENV{MEG_DIR};
my $NEUROMAG_ROOT=$ENV{NEUROMAG_ROOT};
my $dir=$ENV{PWD};


open(ANALYSIS_DISK_FILE, "$MEG_DIR/analysis_disk");
my $analysis_disk = <ANALYSIS_DISK_FILE>;
chomp ($analysis_disk);
close(ANALYSIS_DISK_FILE);


my $script='analysis.sh';			# name of analysis script to be put in each subjects analysis directory
my $btmin=-200;					# default beginning timepoint of baselines in ms
my $btmax=0;					# default end timepoint of baselines in ms


my ($i,$j,$k,$l,$m,$n,$list,$analysis,$type,$sub,$baselines,$tmin,$tmax,$gradonly,$nocov,$sss,$skip,$over,$parent,$bad_file,$st);
my (@conds,@runs,@subjects,@old,@new,@trigs,@sets,@skipfile,@st_file);
my @args=@ARGV;

if ($ARGV[0]){
    while($ARGV[0]){
	my $flag=shift;
	if($flag eq '--runs'){
	    @runs=();
	    while ($ARGV[0]&&$ARGV[0]!~m/^--/){
		$runs[$i++]=shift;
	    }
	}elsif($flag eq '--conds'){
	    @conds=();
	    while ($ARGV[0]&&$ARGV[0]!~m/^--/){
		$conds[$j++]=shift;
	    }
	}elsif($flag eq '--old'){
	    @old=();
	    while ($ARGV[0]&&$ARGV[0]!~m/^--/){
		$old[$k++]=shift;
	    }
	}elsif($flag eq '--new'){
	    @new=();
	    while ($ARGV[0]&&$ARGV[0]!~m/^--/){
		$new[$l++]=shift;
	    }
	    @trigs=@new;
	}elsif($flag eq '--trigs'){
	    @trigs=();
	    while ($ARGV[0]&&$ARGV[0]!~m/^--/){
		$trigs[$m++]=shift;
	    }
	}elsif($flag eq '--sets'){
	    @sets=();
	    while ($ARGV[0]&&$ARGV[0]!~m/^--/){
		$sets[$n++]=shift;
	    }
	}elsif($flag eq '--analysis'){
	    $analysis=shift;
	}elsif($flag eq '--type'){
	    $type=shift;
	}elsif($flag eq '--list'){
	    $list=shift;
	}elsif($flag eq '--over'){
	    $over=1;
	}elsif($flag eq '--gradonly'){
	    $gradonly=1;
	}elsif($flag eq '--baselines'){
	    $baselines=1;
	}elsif($flag eq '--tmin'){
	    $tmin=shift;
	}elsif($flag eq '--tmax'){
	    $tmax=shift;
	}elsif($flag eq '--btmin'){
	    $btmin=shift;
	}elsif($flag eq '--btmax'){
	    $btmax=shift;
	}elsif($flag eq '--nocov'){
	    $nocov=1;
	}elsif($flag eq '--sss'){
	    $sss=1;
	}elsif($flag eq '--st_file'){  # for sss only
	    $st=shift;
	    open(STFILE, "$st");
	    @st_file = <STFILE>; 
	    chomp @st_file;
	}elsif($flag eq '--skip_file'){  # for sss only
	    $skip=shift;
	    open(SKIPFILE, "$skip");
	    @skipfile = <SKIPFILE>; 
	    chomp @skipfile;
	}elsif($flag eq '--parent'){
	    $parent=shift;
	}elsif($flag eq '--bad_file'){
	    $bad_file=shift;
	}else{
	    print "\nInvalid Argument: $flag !\n";
	    exit 1;
	}
    }
}

if ($gradonly && $sss) {
    print 'WARNING: SSS requires using magnetometer and gradiometer channels in order to have enough sources.\n';
}

if($list){
    open (LIST,$list);
    while(my $line=<LIST>){
	chomp($line);
	push @subjects,$line;	
    }
    close (LIST);
}else{
    if(($ENV{PWD}=~m/(mano\S{3})/)||($ENV{PWD}=~m/(average7)/)){
	$sub=$1;
    }else{
	print 'The subject could not be determined';
	exit 1;
    }
    @subjects=$sub;
}

my $batch="batch_${analysis}.sh";		# path to batch script to run analysis scripts for all subjects


$tmin=$tmin/1000;
$tmax=$tmax/1000;

my $btmin_s=$btmin/1000;
my $btmax_s=$btmax/1000;

# (for sss only) machines to run MaxFilter on (randomize order of usage for parallel case)
my @megwsnames = ('megws2','megws3');
@megwsnames = shuffle(@megwsnames);

if($list){
    open (BATCH,">$batch");
    print BATCH "#!/bin/csh -f\n\n";
    print BATCH "foreach sub ( ";
}	

if (! (-d "$analysis_disk/$analysis") ) {
    mkdir "$analysis_disk/$analysis";
}

foreach my $s (@subjects){
    print "$s...";
    if($list){ print BATCH "$s "; }
    if ($s eq 'mano013'){
	@runs=('01','02','03','04','05','06','012');
    }elsif ($s eq 'manoS17'){
	@runs=('01','02','03','04','05','06');
    }else{
	@runs=('01','02','03','04','05','06','012','022');
    }
    chdir "$MEG_DIR/$s";
    if($over){
	if (-d $analysis) {
	    system("rm -Rf $analysis");
	} elsif (-l $analysis) {	    
	    system("rm -Rf " . get_symlink_referent($analysis));
	    system("rm -f " . $analysis);
	}
    }elsif((-d $analysis)||(-l $analysis)){
	print "\nDirectory exists. Use --over to overwrite.\n";
	exit 1;
    }
  # mkdir $analysis,0770;
    mkdir "$analysis_disk/$analysis/$s",0770 or die;
    symlink "$analysis_disk/$analysis/$s", "$analysis";

    chdir $analysis;
    open (SCRIPT,">$script");
    print SCRIPT 			"#!/bin/csh -f\n#\n";
    print SCRIPT 			"# $0 @args\n#\n";
    print SCRIPT 			"# Analysis:		$analysis\n";
    $parent ? print SCRIPT 		"# Parent:		$parent\n" : print SCRIPT 	"# Parent:		-\n";
    print SCRIPT 		        "# Runs:	       	        @runs\n";
    print SCRIPT 			"# Conds:		@conds\n";
    print SCRIPT 			"# Sets:	                @sets\n";
    print SCRIPT 			"# Old:			@old\n";
    print SCRIPT 			"# New:			@new\n";
    print SCRIPT 			"# Trigs:		@trigs\n";
    print SCRIPT 			"# Type:		        $type\n";
    print SCRIPT 			"# tmin:		        $tmin\n";
    print SCRIPT 			"# tmax:		        $tmax\n";
    $list ? print SCRIPT 		"# List:		        $list\n" : print SCRIPT 	"# List:		none\n";
    $nocov ? print SCRIPT 		"# Covariance:		NO\n" : print SCRIPT 		"# Covariance:		YES\n";
    if($baselines){
	print SCRIPT 		"# Baselines:		Calculate\n";
    }elsif($parent){
	print SCRIPT 		"# Baselines:		Copy\n";
    }else{
	print SCRIPT 		"# Baselines:		None\n";
    }
    $over ? print SCRIPT 		"# Overwrote:		YES\n" : print SCRIPT 		"# Overwrote:		NO\n";
    $gradonly ? print SCRIPT 	"# Gradonly:		YES\n" : print SCRIPT 		"# Gradonly:		NO\n";
    $bad_file ? print SCRIPT 	"# Bad File:		$bad_file\n" : print SCRIPT 	"# Bad File:		default\n";
    $sss ? print SCRIPT 	        "# SSS:                  YES\n" : print SCRIPT 	        "# SSS:                  NO\n";
    $st ? print SCRIPT 	        "# Use tSSS subj list:   YES\n" : print SCRIPT 	        "# Use tSSS subj list:   NO\n";
    $skip ? print SCRIPT 	        "# Use skip file:        YES\n" : print SCRIPT 	        "# Use skip file         NO\n";
    
    print SCRIPT "\n\n";
    print SCRIPT "source /space/orsay/8/megdev/megsw/mne/setup/mne/mne_setup_analysis_csh\n";
    print SCRIPT "ln -s \$MEG_DIR/$s/subj_$s/0*/*_raw.fif .\n";
    print SCRIPT "cp \$MEG_DIR/$s/subj_$s/0*/scoring/NEW/*_${type}.lst .\n";
    print SCRIPT "rename _${type}. . *\n";
    print SCRIPT "sm_alias.pl --old @old --new @new --runs @runs\n";
    print SCRIPT "mne_mark_bad_channels ";
    foreach my $sub (@runs){
	print SCRIPT "${s}_${sub}_raw.fif ";
    }
    print SCRIPT "\n";
    $bad_file ? print SCRIPT "bad_channels.pl --bad_file $bad_file" : print SCRIPT "bad_channels.pl";
    
    $gradonly ? print SCRIPT " --gradonly\n" : print SCRIPT "\n";

    print SCRIPT "mne_mark_bad_channels ";
    foreach my $r (@runs){
	print SCRIPT "--bad ${s}.bad ${s}_${r}_raw.fif ";
    }
    print SCRIPT "\n";	
    
    my $i=0;
    if ( $sss ) {
	foreach my $r (@runs) {
	    if ($i >= (scalar(@megwsnames)-1)) { $i = 0; } else { $i++; }
	    my $megwsmachine=$megwsnames[$i];
	    print SCRIPT "/usr/bin/rsh ${megwsmachine} \"cd $MEG_DIR$s/$analysis ; ${NEUROMAG_ROOT}/bin/util/maxfilter -v -format short -autobad off -frame head -origin 0 0 40 ";
	    if (@st_file){ 
		if (grep /$s/, @st_file) {
		    print SCRIPT "-st ";
		}
	    }
	    print SCRIPT "-f ${s}_${r}_raw.fif ";
	    if ($skip) {
		my @skips = grep /^$s $r/, @skipfile; 
		if ( ! ( scalar(@skips) eq 0 ) ) {
		    print SCRIPT "-skip ";
		    foreach my $array_ref (@skips) { 
			my @fields = split ' ', $array_ref;
			print SCRIPT "${fields[2]} ${fields[3]} ";
		    }
		}
	    }
	    print SCRIPT "-o ${s}_${r}_sss.fif \"\n";		  
	}
    }
    
    print SCRIPT "make_ave.pl --conds @conds --trigs @trigs --tmin $tmin --tmax $tmax ";
    
    $nocov ? print SCRIPT "--nocov\n" : print SCRIPT "\n";
    
    print SCRIPT "mne_process_raw ";
    foreach my $r (@runs){
	if ($sss) { print SCRIPT "--raw ${s}_${r}_sss.fif --ave ${s}_${r}.ave "; }
	else {print SCRIPT "--raw ${s}_${r}_raw.fif --ave ${s}_${r}.ave "; }
	unless($nocov){
	    print SCRIPT "--cov ${s}_${r}.cov ";
	}
    }
    if ($sss) { print SCRIPT "--lowpass 30 --lowpassw 5 --projoff --gave ${s}.fif "; }
    else { print SCRIPT "--lowpass 30 --lowpassw 5 --projon --gave ${s}.fif "; }
    $nocov ? print SCRIPT "\n" : print SCRIPT "--gcov ${s}-cov.fif\n";

    if($baselines){
	print SCRIPT "mkdir baselines\n";
	foreach my $r (@runs){
	    print SCRIPT "grep -v omit ${s}_${r}_offl.log > ${s}_${r}-bas.lst\n";
	}
	print SCRIPT "make_ave.pl --conds @conds --trigs @trigs --tmin $btmin_s --tmax $btmax_s --bas\n";
	print SCRIPT "mne_process_raw ";
	foreach my $r (@runs){
	    if ($sss) {
		print SCRIPT "--raw ${s}_${r}_sss.fif --ave ${s}_${r}-bas.ave ";
	    } else {
		print SCRIPT "--raw ${s}_${r}_raw.fif --ave ${s}_${r}-bas.ave ";
	    }
	}
	if ($sss) { print SCRIPT "--lowpass 30 --lowpassw 5 --projoff --gave ${s}-bas.fif\n" }
	else { print SCRIPT "--lowpass 30 --lowpassw 5 --projon --gave ${s}-bas.fif\n" }
	
	for(my $x=0;$x<@conds;$x++){
	    print SCRIPT "mne_change_baselines --in ${s}-bas.fif --bmin $btmin --bmax $btmax --set ${sets[$x]} --list baselines/${s}_${conds[$x]}.bas\n";
	}
	print SCRIPT "mv *-bas* baselines/\n";
    }elsif($parent){
	print SCRIPT "mkdir baselines\n";
	print SCRIPT "cp \$MEG_DIR/${s}/${parent}/baselines/${s}_*.bas baselines/\n";
    }
    
    if($baselines || $parent){
	for(my $x=0;$x<@conds;$x++){
	    print SCRIPT "mne_change_baselines --baselines baselines/${s}_${conds[$x]}.bas --in ${s}.fif --set ${sets[$x]} --out ${s}_${conds[$x]}.fif\n";
	}
    }
    
    unless($nocov){
	foreach my $r (@runs){
	    print SCRIPT "mne_do_forward_solution --meas ${s}_${r}_offl.fif --megonly --overwrite --subject ${s} --fwd ${s}_${r}-fwd.fif\n";
	}
	print SCRIPT "mne_average_forward_solutions";
	foreach my $r (@runs){
	    print SCRIPT " --fwd ${s}_${r}-fwd.fif";
	}
	print SCRIPT " --out ${s}-fwd.fif\n";
	print SCRIPT "mne_do_inverse_operator --meg --loose 0.4 --senscov ${s}-cov.fif --fwd ${s}-fwd.fif --bad ${s}.bad --inv ${s}-inv.fif\n";
    }elsif($parent){
	print SCRIPT "ln -s \$MEG_DIR/${s}/${parent}/${s}-inv.fif .\n";
    }
    print SCRIPT "sm_naves.pl --conds @conds\n";	
    print SCRIPT "rp_scriptlog.sh $MEG_DIR/$s/$analysis/analysis.sh $MEG_DIR/logs/MEG.log";
    print SCRIPT "\n";
    close (SCRIPT);
    chmod 0754,$script;
    print "done!\n";
}

if($list){
    print BATCH ")\n\n";
    print BATCH "cd \$MEG_DIR/\$sub/$analysis\n";
    print BATCH "./$script\n";
    print BATCH "end\n";
    close (BATCH);
    chmod 0754,$batch;
}

sub shuffle {
    my $i=scalar(@_);
    my $j;
    foreach my $item (@_ )
    {
        --$i;
        $j = int rand ($i+1);
        next if $i == $j;
        @_ [$i,$j] = @_[$j,$i];
    }
    return @_;
}

sub get_symlink_referent {
    my $sl=$_[0];
    my $orig=cwd();
    chdir $sl or die;
    my $dest=cwd();
    chdir $orig;
    return $dest;
}