#!/usr/bin/env python
import warnings
warnings.filterwarnings('ignore', '.*negative int.*')
import os, sys, optparse
from sets import Set
import fsutils

# Original Version - Douglas Greve, MGH
# Rewrite - Krish Subramaniam, MGH
# $Id: aparcstats2table,v 1.19 2010/02/17 01:00:58 krish Exp $

# this is a map of measure requested and its corresponding column# in ?h.aparc*.stats
measure_field_map = {'area':2, 'volume':3, 'thickness':4, 'thicknessstd':5, 'meancurv':6 }
# map of delimeter choices and string literals
delimiter2char = {'comma':',', 'tab':'\t', 'space':' ', 'semicolon':';'}

helptext = """
Converts a cortical stats file created by recon-all and or
mris_anatomical_stats (eg, ?h.aparc.stats) into a table in which
each line is a subject and each column is a parcellation. By
default, the values are the area of the parcellation in mm2. The
first row is a list of the parcellation names. The first column is
the subject name.

The subjects list can be specified on either of two ways:
  1. Specify each subject after a -s flag 

            -s subject1 -s subject2 ... --hemi lh
  
  2. Specify all subjects after --subjects flag. --subjects does not have
     to be the last argument. Eg:

            --subjects subject1 subject2 ... --hemi lh

By default, it looks for the ?h.aparc.stats file based on the
Killiany/Desikan parcellation atlas. This can be changed with
'--parc parcellation' where parcellation is the parcellation to
use. An alternative is aparc.a2009s which was developed by
Christophe Destrieux. If this file is not found, it will exit
with an error unless --skip in which case it skips this subject
and moves on to the next.

By default, the area (mm2) of each parcellation is reported. This can
be changed with '--meas measure', where measure can be area, volume
(ie, volume of gray matter), thickness, thicknessstd, or meancurv.
thicknessstd is the standard dev of thickness across space.

Example:
 aparcstats2table --hemi lh --subjects 004 008 --parc aparc.a2009s 
    --meas meancurv -t lh.a2009s.meancurv.txt

lh.a2009s.meancurv.txt will have 3 rows: (1) 'header' with the name
of each structure, (2) mean curvature for each structure for subject

The --common-parcs flag writes only the ROIs which are common to all 
the subjects. Default behavior is it puts 0.0 in the measure of an ROI
which is not present in a subject. 

The --report-rois flag, for each subject, gives what ROIs that are present
in atleast one other subject is absent in current subject and also gives 
what ROIs are unique to the current subject.

The --transpose flag writes the transpose of the table. 
This might be a useful way to see the table when the number of subjects is
relatively less than the number of ROIs.

The --delimiter option controls what character comes between the measures
in the table. Valid options are 'tab' ( default), 'space', 'comma' and
'semicolon'.

The --skip option skips if it can't find a .stats file. Default behavior is
exit the program.

The --parcid-only flag writes only the ROIs name in the 1st row 1st column
of the table. Default is hemi_ROI_measure
"""

def options_parse():
    """
    Command Line Options Parser for aparcstats2table
    initiate the option parser and return the parsed object
    """
    parser = optparse.OptionParser(version='$Id: aparcstats2table,v 1.19 2010/02/17 01:00:58 krish Exp $', usage=helptext)
    
    # help text
    h_sub = '(REQUIRED) subject1 <subject2 subject3..>'
    h_s = ' subjectname'
    h_hemi = '(REQUIRED) lh or rh'
    h_parc = 'parcellation.. default is aparc ( alt aparc.a2009s)'
    h_meas = 'measure: default is area ( alt volume, thickness, thicknessstd, meancurv)'
    h_skip = 'if a subject does not have input, skip it'
    h_t = '(REQUIRED) output table file'
    h_deli = 'delimiter between measures in the table. default is tab (alt comma, space, semicolon )' 
    h_parcid = 'do not pre/append hemi/meas to parcellation name'
    h_common = 'output only the common parcellations of all the subjects given'
    h_roi = 'print ROIs information for each subject'
    h_tr = 'transpose the table ( default is subjects in rows and ROIs in cols)' 
    h_v = 'increase verbosity'

    # Add the options
    parser.add_option('--subjects', dest='subjects' ,action='callback',
                      callback=fsutils.callback_var,  help=h_sub)
    parser.add_option('-s', dest='subjects' ,action='append',
                      help=h_s)
    parser.add_option('--hemi', dest='hemi',
                      choices=('lh','rh'), help=h_hemi)
    parser.add_option('-t', '--tablefile', dest='outputfile',
                      help=h_t)
    parser.add_option('-p', '--parc', dest='parc',
                      default='aparc', help=h_parc)
    parser.add_option('-m', '--measure', dest='meas',
                      choices=('area','volume','thickness','thicknessstd','meancurv'),
                      default='area', help=h_meas)
    parser.add_option('-d', '--delimiter', dest='delimiter',
                      choices=('comma','tab','space','semicolon'),
                      default='tab', help=h_deli)
    parser.add_option('--skip', action='store_true', dest='skipflag',
                      default=False, help=h_skip)
    parser.add_option('--parcid-only', action='store_true', dest='parcidflag',
                      default=False, help=h_parcid)
    parser.add_option('--common-parcs', action='store_true', dest='commonparcflag',
                      default=False, help=h_common)
    parser.add_option('--report-rois', action='store_true', dest='reportroiflag',
                      default=False, help=h_roi)
    parser.add_option('', '--transpose', action='store_true', dest='transposeflag',
                      default=False, help=h_tr)
    parser.add_option('-v', '--debug', action='store_true', dest='verboseflag',
                      default=False, help=h_v)
    (options, args) = parser.parse_args()
   
    # error check
    if options.subjects is not None:
        if len(options.subjects) < 1:
            print 'ERROR: atleast 1 subject must be provided'
            sys.exit(1)
    else:
        print 'ERROR: --subjects sub1 .. is mandatory. See --help for more info'
        sys.exit(1)

    if not options.outputfile:
        print 'ERROR: output table name should be specified'
        sys.exit(1)
    if not options.hemi:
        print 'ERROR: hemisphere should be provided (lh or rh)'
        sys.exit(1)

    return options

def roi_setops(table, subjects):
    """
    This outputs a tuple which has 2 lists
    One which is the union of all ROIs of all subjects
    And another which is a list of ROIs only common to all subjects
    """
    _all_parcs = []
    for subject in subjects:
        _all_parcs.append( table[subject].keys())
    # we have a list of lists. Flatten it. 
    temp_parcs = [item for sublist in _all_parcs for item in sublist]
    all_parcs = fsutils.unique_union(temp_parcs)
    
    common_parcs = table[subjects[0]].keys()
    for subject in subjects[1:]:
        common_parcs = fsutils.intersect_order(common_parcs, table[subject].keys())
    
    return (all_parcs, common_parcs) 

def build_table(options):
    """
    This function builds the 2d table (spreadsheet )
    of subjects vs ROIs. The values are given by what the user chose in 
    measure option.
    """
    # check the subjects dir
    subjdir = fsutils.check_subjdirs()
    
    # init the table
    table = fsutils.Ddict(fsutils.StableDict)

    o = options
    if o.verboseflag: print '\nParsing individual stats files\n' + '-'*40 
    # for every subject
    for subject in o.subjects:
        statsfile = os.path.join(subjdir, subject, 'stats', o.hemi + '.' + o.parc + '.stats')
        if o.verboseflag: 
            print 'Subject:%s' %subject 
            print 'Stats File:%s' %statsfile,
        if os.path.exists(statsfile):
            # if file size is less than 10, not a valid parcstats file. so exit
            if os.path.getsize(statsfile) < 10:
                print 'WARNING: File:' + statsfile + ' is not a valid aparc statsfile'
                print 'Exclude it from the list of subjects/don\'t use common-segs option'
            fp = open(statsfile, 'r')
            if o.verboseflag: print '\nParcellations: ',
            for line in fp:
                # a valid line is a line without a '#'
                if line.rfind('#') == -1:
                    strlist = line.split()
                    # for every parcellation
                    parcid = strlist[0]
                    if o.verboseflag: print parcid + ' ',
                    val = float(strlist[measure_field_map[o.meas]])
                    # assign the measure to the 2d table
                    table[subject][parcid] = val
        else:
            if not o.skipflag:
                print 'ERROR: cannot find %s' %statsfile
                print 'Use --skip flag if you want to continue in such cases'
                sys.exit(1)
            else:
                print 'WARNING: cannot find %s' %statsfile
                print 'skipping the subject %s' %subject
        if o.verboseflag: print '\n'
   
    (all_parcs, common_parcs) = roi_setops(table, o.subjects)
    if not o.commonparcflag:
        # The ROIs might not be the same in every subject. But mostly they overlap
        # What we need is the UNION of all the ROI labels of all subjects. 
        # The result is if 'bert' doesn't have the ROI 'putamen' 
        # table['bert']['putamen'] doesn't exist. We
        # should make it 0 somehow. Thus we get a nice grid which we can output.
        
        new_table = fsutils.Ddict(fsutils.StableDict)
        for subject in o.subjects:
            for parc in all_parcs:
                if not parc in table[subject]:
                    new_table[subject][parc] = 0.0
                else:
                    new_table[subject][parc] = table[subject][parc]
        cols = all_parcs
    else:
        # common parc is enabled. ROIs which are common to ALL the subjects are
        # only outputted. So we should take sets of all ROI labels of all
        # subjects and intersect them
        # Note that for union, we init'ed with an empty set.
        # Here we init the intersection with the ROIs for the first subject 

        new_table = fsutils.Ddict(fsutils.StableDict)
        for subject in o.subjects:
            for parc in common_parcs:
                new_table[subject][parc] = table[subject][parc]
        cols = common_parcs

    if o.reportroiflag:
        print '\n\nROIs information'
        print  '-'*40
        print 'ROIs which are found in atleast one subject ( union ): ',
        for parc in all_parcs:
            print parc+' ',
        print '\n\nROIs which are common to all subjects ( intersection ): ',
        for parc in common_parcs:
            print parc+' ',
        print '\n\nROIs information pertaining to each subject'
        print '-'*40
        for s in o.subjects:
            roi_set_of_s = Set(table[s].keys())
            print 'Subject: %s' %s
            print 'ROIs which are not found in this subject but found in alteast one other subject: ',
            tmpset = all_parcs - roi_set_of_s
            for parc in tmpset:
                print parc+' ',
            print '\nROIs which are found only in this subject: ',
            tmpset = roi_set_of_s - common_parcs 
            for parc in tmpset:
                print parc+' ',
            print '\n'

    return (o.subjects, list(cols), new_table)
        
def write_table(options, rows, cols, table):
    """
    Write the table from memory to disk. Initialize the writer class.
    """
    tw = fsutils.TableWriter(rows, cols, table)
    r1c1 = '%s.%s.%s' %(options.hemi, options.parc, options.meas)
    tw.assign_attributes(filename=options.outputfile, row1col1=r1c1,
                         delimiter=delimiter2char[options.delimiter] )
    # we might need the hemisphere and measure info in columns as well 
    if not options.parcidflag:
        tw.decorate_col_titles(options.hemi+'_', '_'+options.meas)
    if options.transposeflag:
        tw.write_transpose()
    else:
        tw.write()
        
  
if __name__=="__main__":
    options = options_parse()
    print 'Number of subjects : %d' %len(options.subjects)
    print 'Building the table..'
    (rows, cols, table) = build_table(options)
    print '\nWriting the table to %s' %options.outputfile
    write_table(options, rows, cols, table)
    sys.exit(0) 
