#!/Users/boydb1/anaconda/bin/python
# -*- coding: utf-8 -*-

'''
Generate a report for a list of project on XNAT

@author: Benjamin Yvernault, Electrical Engineering, Vanderbilt University
'''
import os
import sys
import logging
from dax import XnatUtils
from datetime import datetime

DEFAULT_CSV_LIST=['object_type','project_id','subject_label','session_type','session_label', 'as_label', 'as_type', 'as_description','quality','resource']
ORDER_PRINTING=['commun','project','subject','session','scan','assessor','resource']
VARLIST=['scan_id', 'type', 'series_description', 'quality', 'note', 'frames', 'resource', 'assessor_id', 'assessor_label', 'assessor_URI', 'proctype', 'procstatus', 'qcstatus', 'version', 'jobid', 'memused', 'walltimeused', 'jobnode', 'jobstartdate', 'object_type', 'project_id', 'subject_id', 'subject_label', 'session_id', 'session_type', 'session_label', 'handedness', 'gender', 'yob', 'age', 'last_modified', 'last_updated']
VARIABLES_LIST={'commun'   : ['object_type'],
                'project'  : ['project_id'],
                'subject'  : ['subject_id','subject_label','handedness','gender','yob'],
                'session'  : ['session_id','session_type','session_label','age','last_modified','last_updated'],
                'scan'     : ['scan_id','type','series_description','quality','note','frames'],
                'assessor' : ['assessor_id','assessor_label','assessor_URI','proctype','procstatus','qcstatus','version','jobid','memused','walltimeused','jobnode','jobstartdate'],
                'resource' : ['resource']}

########################################## USEFUL FUNCTIONS ##########################################
def setup_info_logger(name,logfile):
    if logfile:
        handler=logging.FileHandler(logfile,'w')
    else:
        handler=logging.StreamHandler()
    
    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)
    logger.addHandler(handler)
    return logger
    
########################################## SPECIFIC FUNCTIONS ##########################################
def report(xnat,options,project_list):
    logger.info('Date: '+str(datetime.now())+'\n')
    #print header:
    if not options.format:
        logger.info(','.join(DEFAULT_CSV_LIST))
    else:
        logger.info(options.format)
    for project in project_list:
        if not options.format:
            default_report(xnat,project)
        else:
            header=options.format.split(',')
            customize_report(xnat,project,header)
    
def default_report(xnat,project):
    #get list from XNAT
    scans_list=XnatUtils.list_project_scans(xnat, project)
    assessors_list=XnatUtils.list_project_assessors(xnat,project)
    #get list of subject
    subjects_list=set([s['subject_label'] for s in scans_list])
    #Loop through subjects / loop through scan/assessor if needed
    for subject in subjects_list:
        #SCAN
        for scan in [s for s in scans_list if s['subject_label']==subject]:
            SResources='/'.join([r['label'] for r in XnatUtils.list_scan_resources(xnat,project,scan['subject_label'],scan['session_label'],scan['ID'])])
            logger.info(','.join(['scan',scan['subject_label'],scan['session_type'],scan['session_label'],scan['ID'],scan['type'],scan['series_description'],scan['quality'],SResources]))                  
        #ASSESSOR
        for assessor in [a for a in assessors_list if a['subject_label']==subject]:
            AResources='/'.join([r['label'] for r in XnatUtils.list_assessor_out_resources(xnat,project,assessor['subject_label'],assessor['session_label'],assessor['label'])])
            logger.info(','.join(['assessor',assessor['subject_label'],assessor['session_type'],assessor['session_label'],assessor['label'],assessor['proctype'],assessor['procstatus'],assessor['qcstatus'],AResources]))
    
def customize_report(xnat,project,header):
    #Loop through subjects / loop through scan/assessor if needed
    if filter(lambda x: x in header,VARIABLES_LIST['scan']) or filter(lambda x: x in header,VARIABLES_LIST['assessor']):
        customize_report_under_sessions(xnat,project,header)
    elif filter(lambda x: x in header,VARIABLES_LIST['session']):
        customize_report_sessions(xnat,project,header)
    elif filter(lambda x: x in header,VARIABLES_LIST['subject']):
        customize_report_subjects(xnat,project,header)
    else:
        logger.info(','.join(get_row(xnat,{'project_id':project},header)))

def customize_report_subjects(xnat,project,header): 
    subjects_list=XnatUtils.list_subjects(xnat,project)
    for subject in subjects_list:
        logger.info(','.join(get_row(xnat,subject,header))) 
        
def customize_report_sessions(xnat,project,header): 
    sessions_list=XnatUtils.list_sessions(xnat,project)
    for session in sorted(sessions_list, key=lambda k: k['session_label']):
        logger.info(','.join(get_row(xnat,session,header)))    
            
def customize_report_under_sessions(xnat,project,header):
    if filter(lambda x: x in header,VARIABLES_LIST['scan']):
        customize_report_scans(xnat,project,header)
    if filter(lambda x: x in header,VARIABLES_LIST['assessor']):
        customize_report_assessors(xnat,project,header)
    
def customize_report_scans(xnat,project,header):
    scans_list=XnatUtils.list_project_scans(xnat, project)
    for scan in sorted(scans_list, key=lambda k: k['subject_label']):
        logger.info(','.join(get_row(xnat,scan,header)))

def customize_report_assessors(xnat,project,header):
    assessors_list=XnatUtils.list_project_assessors(xnat,project)            
    for assessor in sorted(assessors_list, key=lambda k: k['subject_label']):
        logger.info(','.join(get_row(xnat,assessor,header)))

def get_row(xnat,obj_dict,header):
    row=list()
    for field in header:
        if field=='object_type':
            if 'scan_id' in obj_dict.keys():
                row.append('scan')
            elif 'assessor_label' in obj_dict.keys():
                row.append('assessor')
            elif 'session_label' in obj_dict.keys():
                row.append('session')
            elif 'subject_label' in obj_dict.keys():
                row.append('subject')
            else:
                row.append('project')
        elif field=='resource':
            if 'scan_id' in obj_dict.keys():
                Resources='/'.join([r['label'] for r in XnatUtils.list_scan_resources(xnat,obj_dict['project_id'],obj_dict['subject_label'],obj_dict['session_label'],obj_dict['ID'])])
                row.append(str(Resources))
            elif 'assessor_label' in obj_dict.keys():
                Resources='/'.join([r['label'] for r in XnatUtils.list_assessor_out_resources(xnat,obj_dict['project_id'],obj_dict['subject_label'],obj_dict['session_label'],obj_dict['label'])])
                row.append(str(Resources))
            elif 'session_label' in obj_dict.keys():
                Resources='/'.join([r['label'] for r in XnatUtils.list_experiment_resources(xnat,obj_dict['project_id'],obj_dict['subject_label'],obj_dict['session_label'])])
                row.append(str(Resources)) 
            elif 'subject_label' in obj_dict.keys():
                Resources='/'.join([r['label'] for r in XnatUtils.list_subject_resources(xnat,obj_dict['project_id'],obj_dict['subject_label'])])
                row.append(str(Resources))
            elif 'project_id' in obj_dict.keys():
                Resources='/'.join([r['label'] for r in XnatUtils.list_project_resources(xnat,obj_dict['project_id'])])
                row.append(str(Resources))  
        else:
            row.append(obj_dict.get(field)) 
    return row

########################################## CHECK OPTIONS ##########################################
def check_options(options):
    #check format:
    if options.format and options.format[-1]==',':
        options.format=options.format[:-1]
    if options.format and filter(lambda x: x not in VARLIST, options.format.split(',')):
        print 'OPTION ERROR: some variables chosen in format does not exist on XNAT. See below for the available variables names:'
        for key in ORDER_PRINTING:
            print key+' variables:'
            for value in VARIABLES_LIST[key]:
                print ' * %*s ' % (-30,value)
        return False
    if not options.printformat:
        #check options :
        if not options.projects:
            print 'OPTION ERROR: No project selected.Please specify one or more project with option -p/--project.'
            return False
        #check the name given
        if options.csvfile:
            folder=os.path.dirname(os.path.abspath(options.csvfile))
            if not os.path.exists(folder):
                print 'OPTION ERROR: the txt file path <'+folder+'> does not exist. Please check the path given.'
                return False
    return True

########################################## MAIN DISPLAY ##########################################
def Main_display(parser):
    (options,args) = parser.parse_args()
    #Display:
    print '################################################################'
    print '#                          XNATREPORT                          #'
    print '#                                                              #'
    print '# Developed by the masiLab Vanderbilt University, TN, USA.     #'
    print '# If issues, email benjamin.c.yvernault@vanderbilt.edu         #'
    print '# Usage:                                                       #'
    print '#     Print a detailed report from XNAT projects               #'
    print '# Parameters :                                                 #'
    if options=={'csvfile': None,'projects': None,'printformat':False,'format':None}:
        print '#     No Arguments given                                       #'
        print '#     Use "Xnatreport -h" to see the options                   #'
        print '################################################################'
        parser.print_help()
        sys.exit()
    else:
        if options.printformat:
            print '#     %*s -> %*s#' %(-20,'Print variables',-33,'on')
        else:
            if options.projects:
                print '#     %*s -> %*s#' %(-20,'Project(s)',-33,get_proper_str(options.projects))
            if options.format:
                print '#     %*s -> %*s#' %(-20,'Format',-33,get_proper_str(options.format))
            if options.csvfile:
                print '#     %*s -> %*s#' %(-20,'CSV filename',-33,get_proper_str(options.csvfile,True))
        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+="   *Create a report about Xnat projects.\n"
    #Example
    usage+="Examples:\n"
    usage+="   *Report of a project: Xnatreport -p PID \n"
    usage+="   *Report with a specific format: Xnatreport -p PID --format object_type,session_id,session_label,age \n"
    usage+="   *print the format available: Xnatreport --printformat \n"
    usage+="   *Save report in a csv: Xnatreport -p PID -c report.csv\n"
    return usage
    
def parse_args():
    from optparse import OptionParser
    usage = get_usage()
    parser = OptionParser(usage=usage)
    parser.add_option("-p", "--project", dest="projects",default=None,
                  help="List of project ID on Xnat separate by a coma", metavar="PROJECT_ID")
    parser.add_option("-c", "--csvfile", dest="csvfile",default=None,
                  help="csv fullpath where to save the report.", metavar="CSV_FULLPATH_FILE")
    parser.add_option("--format", dest="format",default=None,
                  help="Header for the csv. format: list of variables name separated by a comma.", metavar="COMMA_SEPARATED_LIST")
    parser.add_option("--printformat", dest="printformat",action="store_true",default=False,
                  help="Print available variables names for the option --format.", metavar="COMMA_SEPARATED_LIST")
    return parser
    
if __name__ == '__main__':
    parser=parse_args()
    (options,args) = parser.parse_args()
    #############################
    #Main display:
    Main_display(parser)
    #check options:
    run=check_options(options)
    #############################
    logger = setup_info_logger('xnatreport',options.csvfile)
    
    #############################
    # RUN                       #
    #############################
    if run:
        print '=========================================================================='
        if options.printformat:
            logger.info('INFO: Printing the variables available: ')
            for key in ORDER_PRINTING:
                logger.info(key+' variables:')
                for value in VARIABLES_LIST[key]:
                    logger.info(' * %*s ' % (-30,value))
        else:
            if options.format and options.format[-1]==',':
                options.format=options.format[:-1]
            #Arguments :
            project_list=options.projects.split(',')  
            
            """ Report on Xnat database for the list of project for processes or resources on scan """
            if options.format and options.format=='object_type':
                logger.info('object_type')
                logger.info('project')
                logger.info('subject')
                logger.info('session')
                logger.info('scan')
                logger.info('assessor')
            else:
                if options.format and options.format=='resource':
                    print 'WARNING: you gave only the resource to --formats options. Going to add object_type and project_id.'
                    options.format='object_type,project_id,resource'
                #Starting Report
                try:
                    #connection to Xnat
                    xnat = XnatUtils.get_interface()
                    print "Report for the following project(s):"
                    print '------------------------------------'
                    for project in project_list:
                        #check if the project exists:
                        project_xnat=xnat.select('/project/'+project)
                        if not project_xnat.exists():
                            print ' - ERROR: Project '+project+' does not exist on Xnat.'
                            project_list.remove(project)
                            pass
                        else:
                            print '  -'+project
                    print '------------------------------------'
                    print 'WARNING: extracting information from XNAT for a full project might take some time. Please be patient.\n'
                    #Writing report
                    report(xnat,options,project_list)
                
                finally:                                        
                    xnat.disconnect()
        print '==========================================================================\n'
