#!/Users/boydb1/anaconda/bin/python
# -*- coding: utf-8 -*-
'''
Created on May 6, 2013
Edited on February 19, 2015
@author: Benjamin Yvernault, Electrical Engineering, Vanderbilt University
'''

import os,sys,datetime,csv
from dax import XnatUtils

REQUIRED_VARIABLES_LIST=['project_id','subject_label','session_label']
SUBJECT_PARAMETERS_LIST=['handedness','gender','yob']
SESSION_PARAMETERS_LIST=['age','scanner','scanner_manufacturer','scanner_model','acquisition_site']

########################################## USEFUL FUNCTIONS ##########################################
def get_xnat_object(xnat,obj_dict):
    return xnat.select('/project/'+obj_dict['project_id']+'/subject/'+obj_dict['subject_label'])
    
def get_gender_from_label(gender):
    if gender.lower() in ['female','f']:
        return 'female'
    elif gender.lower() in ['male','m']:
        return 'male'
    else:
        return 'unknown'
    
def get_handedness_from_label(handedness):
    if handedness.lower() in ['right','r']:
        return 'right'
    elif handedness.lower() in ['left','l']:
        return 'left'
    elif handedness.lower() in ['ambidextrous','a']:
        return 'ambidextrous'
    else:
        return 'unknown'
    
def get_yob_from_label(yob,row_index):
    if '/' in yob:
        print ' warning: row '+str(row_index+1)+' -- "/" detected in the year of birth (yob) value. Using last digit of the string: '+yob.split('/')[-1]
        yob=yob.split('/')[-1]
    try:
        _=int(yob)
        return yob
    except:
        print ' warning: row '+str(row_index+1)+' -- year of birth (yob) not a proper integer.'
        return ''
    
def read_csv(options):
    demo_list=list()
    header=list()
    session_variables=list()
    #Get the sessions variables:
    if options.session_format:
        session_variables=set(options.session_format.split(',')+SESSION_PARAMETERS_LIST)
    else:
        session_variables=SESSION_PARAMETERS_LIST
    print 'Warning: the variables for a session are the following : '+','.join(session_variables)
    #Header if set:
    if options.format:
        print 'Warning: the header has been set by the options --format.'
        header=options.format.split(',')
    else:
        print 'Warning: the header has not been set. Using the first line for the header.'
    print 'INFO: Reading CSV -- if you see any warning, it means that the information on the row will not be uploaded.'
    with open(options.csvfile,'rb') as csvfileread:
        csvreader = csv.reader(csvfileread,delimiter=options.delimiter)
        for index,row in enumerate(csvreader):
            if index==0: #possible header, check it out
                if not header:
                    header=row
                else:
                    demo_list.append(get_row_csv(row,header,index,session_variables))
                #Check the header
                if 'project_id' not in header or 'subject_label' not in header:
                    print ' warning: row '+str(index+1)+' -- "project_id" and "subject_label" not in the header. Check the header.'
                    return None,None
                if 'age' in header and not 'session_label' in header:
                    print ' warning: row '+str(index+1)+' -- age found in header but session_label not in it. Check the header.'
            else:
                demo_list.append(get_row_csv(row,header,index,session_variables))
    #Remove empty dictionaries:
    return [x for x in demo_list if x],header

def get_row_csv(row,header,row_index,session_variables):
    tmp_dict=dict()
    if len(row)!=len(header):
        print ' WARNING: row '+str(row_index+1)+' -- less or more columns than the header. Check the csv file.'
        return tmp_dict
    else:
        tmp_dict['upload_demo_subject']=dict()
        tmp_dict['upload_demo_session']=dict()
        for index,value in enumerate(row):
            if header[index] in REQUIRED_VARIABLES_LIST:
                tmp_dict[header[index]]=value
            else:
                if value:
                    if header[index] in session_variables:
                        tmp_dict['upload_demo_session'][header[index]]=value
                    else:
                        if header[index]=='handedness':
                            tmp_dict['upload_demo_subject'][header[index]]=get_handedness_from_label(value)
                        elif header[index]=='gender':
                            tmp_dict['upload_demo_subject'][header[index]]=get_gender_from_label(value)
                        elif header[index]=='yob':
                            tmp_dict['upload_demo_subject'][header[index]]=get_yob_from_label(value,row_index)
                        else:
                            tmp_dict['upload_demo_subject'][header[index]]=value
        return tmp_dict

def smaller_str(str_option,size=10,end=False):
    if len(str_option)>size:
        if end:
            return '...'+str_option[-size:]
        else:
            return str_option[:size]+'...'
    else:
        return str_option
    
def print_report(demo_list,header):
    #Display:
    print 'Report information :'
    print 'List of the data that need to be upload for the header: '
    print ' , '.join(header)
    print '-----------------------------------------------------------------------------'
    newlist = sorted(demo_list, key=lambda k: k['project_id']) 
    for obj_dict in sorted(newlist):
        if obj_dict:
            row=get_row_info(obj_dict,header)
            print ' , '.join(row)
    print '-----------------------------------------------------------------------------\n'

def get_row_info(obj_dict,header):
    row=list()
    for field in header:
        if field in obj_dict.keys():
            row.append(obj_dict.get(field,''))
        else:
            if field in obj_dict['upload_demo_subject'].keys():
                row.append(obj_dict['upload_demo_subject'].get(field,''))
            elif field in obj_dict['upload_demo_session'].keys():
                row.append(obj_dict['upload_demo_session'].get(field,''))
            else:
                row.append('')
    return row

def print_info_row(object_dict):
    print ' *Project: '+object_dict['project_id']+' / Subject: '+object_dict['subject_label'],
    print ' / Session: '+object_dict['session_label'] if 'session_label' in object_dict.keys() and object_dict['session_label'] else '',
    if 'upload_demo_subject' in object_dict.keys():
        for tag,value in object_dict['upload_demo_subject'].items():
            print ' / '+tag+': '+value,
    if 'upload_demo_session' in object_dict.keys():
        for tag,value in object_dict['upload_demo_session'].items():
            print ' / '+tag+': '+value,
    print ''

########################################## UPLOAD FUNCTIONS ##########################################
def upload_demographic_data(demo_list):
    for obj_dict in sorted(demo_list, key=lambda k: k['project_id']):
        print_info_row(obj_dict)
        subject_obj=get_xnat_object(xnat,obj_dict)
        if not subject_obj.exists():
            print " --> WARNING: Subject "+obj_dict['subject_label']+" doesn't exist. No information will be uploaded."
        else:
            update_demo_data(subject_obj,obj_dict)

def update_demo_data(subject_obj,obj_dict):
    if 'upload_demo_subject' in obj_dict.keys():
        for tag,value in obj_dict['upload_demo_subject'].items():
            update_tag_subject(subject_obj,obj_dict['subject_label'],tag,value)
    if 'upload_demo_session' in obj_dict.keys():
        session_obj=subject_obj.experiment(obj_dict['session_label'])
        for tag,value in obj_dict['upload_demo_session'].items():
            update_tag_session(session_obj,obj_dict['session_label'],tag,value)

def update_tag_subject(subject_obj,subject_label,tag,value):
    if tag and value:
        if tag in SUBJECT_PARAMETERS_LIST: # default attributes
            subject_obj.attrs.set('xnat:subjectData/demographics[@xsi:type=xnat:demographicData]/'+tag.lower(),value)
            print "  - "+tag+" set."
        else: # specific attributes
            try: 
                subject_obj.attrs.set("xnat:subjectData/fields/field[name="+tag.lower()+"]/field",value)
                print "  - "+tag+" set."
            except:
                print "  --> warning: tag "+tag+" -- Subject "+subject_label+" doesn't have this tag or wrong value given."
    
def update_tag_session(session_obj,session_label,tag,value):
    if tag and value:
        if not session_obj.exists():
            print "  --> warning: tag "+tag+" -- Session "+session_label+" doesn't exist or not set."
        else:
            if tag in SESSION_PARAMETERS_LIST: # default attributes
                session_obj.attrs.set(session_obj.attrs.get('xsiType')+'/'+tag.lower(),value)
                print "  - "+tag+" set."
            else: # specific attributes
                try:
                    session_obj.attrs.set(session_obj.attrs.get('xsiType')+"/fields/field[name="+tag.lower()+"]/field",value)
                    print "  - "+tag+" set."
                except:
                    print "  --> warning: tag "+tag+" -- Session "+session_label+" doesn't have this tag."

########################################## MAIN DISPLAY ##########################################
def Main_display(parser):
    (options,args) = parser.parse_args()
    print '################################################################'
    print '#                         XNATDEMOGRAPHIC                      #'
    print '#                                                              #'
    print '# Developed by the masiLab Vanderbilt University, TN, USA.     #'
    print '# If issues, email benjamin.c.yvernault@vanderbilt.edu         #'
    print '# Usage:                                                       #'
    print '#     Upload demographic data to XNAT from a csv file          #'
    print '# Parameters :                                                 #'
    if options=={ 'csvfile':None,'delimiter': ',', 'report': False,'printformat':False,'format':None}:
        print '#     No Arguments given                                       #'
        print '#     Use "XnatDemographic -h" to see the options              #'
        print '################################################################'
        parser.print_help()
        sys.exit()
    else:
        if options.printformat:
            print '#     %*s -> %*s#' %(-20,'Print format',-33,'on') 
        else:
            if options.csvfile:
                print '#     %*s -> %*s#' %(-20,'Input File Path',-33,get_proper_str(options.csvfile,True))
            if options.delimiter:
                print '#     %*s -> %*s#' %(-20,'Delimiter for output',-33,get_proper_str(options.delimiter))
            if options.format:
                print '#     %*s -> %*s#' %(-20,'Header format',-33,get_proper_str(options.format))
            if options.session_format:
                print '#     %*s -> %*s#' %(-20,'Session attributes',-33,get_proper_str(options.session_format))
            if options.report:
                print '#     %*s -> %*s#' %(-20,'Report Mode',-33,'on')   
        print '################################################################'
        
def get_proper_str(str_option,end=False):
    if len(str_option)>32:
        if end:
            return '...'+str_option[-29:]
        else:
            return str_option[:29]+'...'
    else:
        return str_option

def get_usage():
    usage="usage: %prog [options] \n"
    usage+="What is the script doing : \n"
    usage+="   * Upload demographic data to a project on XNAT from a CSV file giving to the script as an input.\n"
    usage+="     Demographic data can only be upload on a SUBJECT or SESSION with this tool.\n"
    usage+="     Default parameters that can be uploaded to XNAT: handedness/gender/yob/age/scanner/scanner_manufacturer/scanner_model/acquisition_site'.\n"
    usage+="     You can also upload your specific parameters but you need to create them first on XNAT.\n"
    usage+="     By default, specific parameters are supposed to be part of the subject level.\n\n"
    #Example
    usage+="Examples:\n"
    usage+="   *Upload demographic data handedness/gender/age: XnatDemographic -i demographic_data.csv \n"
    usage+="   *See the report before uploading: Xnatdemographic -i demographic_data.csv \n"
    usage+="   *Upload demographic data with a file delimited by a /: Xnatdemographic -i demographic_data.csv --delimiter=/ \n"
    usage+="   *Upload demographic data with a different header that the default one: Xnatdemographic -i demographic_data.csv --format=project_id,subject_label,session_label,race,handedness,gender,age\n"
    return usage

########################################## OPTIONS ##########################################
def parse_args():
    from optparse import OptionParser
    usage = get_usage()
    parser = OptionParser(usage=usage)
    #need this options
    parser.add_option("-c","--csv",dest="csvfile",default=None,
                  help="File path as inputs that will be read for XNAT information. Default header to use: project_id,subject_label,session_label,handedness,gender,age.", metavar="FILEPATH")
    parser.add_option("-d","--delimiter",dest="delimiter",default=',',
                  help="Delimiter for the output file. Default: comma.", metavar="STRING")
    parser.add_option("--format", dest="format",default=None,
                  help="Header for the csv. format: list of variables name separated by a comma. Default: using first line in the csv for the header.", metavar="COMMA_SEPARATED_LIST")
    parser.add_option("--sessformat",dest="session_format",default=None,
                  help="List of specific attributes for a session. format: comma separated list.", metavar="COMMA_SEPARATED_LIST")
    parser.add_option("--report",dest="report",action="store_true", default=False,
                  help="Show what information the script will upload to XNAT.", metavar="FILEPATH")
    parser.add_option("--printformat",dest="printformat",action='store_true',default=False,
                  help="Print available parameters for the csv header.", metavar="")
    return parser

########################################## MAIN FUNCTION ##########################################
if __name__ == '__main__':
    parser=parse_args()
    (options,args) = parser.parse_args()
    #############################
    #Main display:
    Main_display(parser)
    #############################
    # RUN                       #
    #############################
    print '=========================================================================='
    if options.printformat:
        print 'INFO: Printing the variables available for formating the output: '
        print 'Required variables:'
        for value in REQUIRED_VARIABLES_LIST:
            print ' * %*s ' % (-30,value)
        print 'Default Subject demographic attributes:' 
        for value in SUBJECT_PARAMETERS_LIST:
            print ' * %*s ' % (-30,value)
        print 'Default Session demographic attributes:' 
        for value in SESSION_PARAMETERS_LIST:
            print ' * %*s ' % (-30,value)
        print'\nINFO: You can generate your own attributes by creating them on the XNAT directly and use the script to upload the values.'
        print'WARNING:'
        print'    * you need to have "project_id,subject_label" if you upload demographic data to a subject.'
        print'    * for session attributes, you need to give the session_label as well.'
        print'    * for specific attributes that you created on XNAT, please put the name of the variables in the header.'
    if options.csvfile:
        if not os.path.exists(options.csvfile):
            print 'ERROR: the csv file '+options.csvfile+' does not exist.'
        else:
            demo_list,header=read_csv(options)
            if not demo_list:
                print 'INFO: no valid demographic found.'
            elif options.report:
                print_report(demo_list,header)
            else:      
                try:
                    xnat=XnatUtils.get_interface()
                    print 'INFO: Uploading demographics data to XNAT. It will take some time to upload demographic data. Please be patient.'
                    upload_demographic_data(demo_list)
                finally:
                    xnat.disconnect()
    print '=========================================================================='
