#!/Users/boydb1/anaconda/bin/python
# -*- coding: utf-8 -*-
'''
Created on Jan 24, 2013
Upgraded on February 19, 2015

@author: yvernabc
'''

import os,sys
from dax import XnatUtils

########################################## USEFUL FUNCTIONS ##########################################
def get_option_list(option):
    if not option:
        return None
    elif option=='all':
        return 'all'
    elif option=='nan':
        return None
    else:
        return option.split(',')

def isFreeSurfer(proctype):
    return proctype=='FS' or proctype=='FreeSurfer'
                    
def get_assessor_dict(assessor_label):
    assessor_dict=dict()
    labels=assessor_label.split('-x-')
    if len(labels)==1:
        print'ERROR: WRONG PROCESS LABEL: the assessor label can not be set (ERROR no "-x-" in the name)'
        print'  -> Skipping the processor '+assessor_label
    else:
        assessor_dict['project_id']=labels[0]
        assessor_dict['subject_label']=labels[1]
        assessor_dict['session_label']=labels[2]
        assessor_dict['label']=assessor_label
        assessor_dict['proctype']=labels[-1]
    return assessor_dict

########################################## XNAT FUNCTIONS ##########################################
def xnat_list_assessors(xnat,options):
    projects_assessors_list=list()
    print 'INFO: Querying XNAT to get assessors labels for all the projects.'
    for project in get_option_list(options.project):
        sys.stdout.write("  *Project: "+project+'\t\t\t\t\t\t\n')
        #for all processors in the project:
        proc_list=XnatUtils.list_project_assessors(xnat, project)
        if not proc_list:
            sys.stdout.write("   !!ERROR: You don't have access to the project: "+project+".!!\n")
            continue
        #Filters for subjects/sessions/Proctypes/status
        projects_assessors_list.extend(filter_assessors(options,proc_list))
    
    #Print number of assessors found:
    print 'INFO: Number of XNAT assessors found after filters for all the project:'
    if projects_assessors_list:
        print ' ---------------------------------------------'
        print '| %*s | %*s |' % (-20,'Project ID',-20,'Number of Assessors')
        print ' ---------------------------------------------'
        for project in set([o['project_id'] for o in projects_assessors_list]):
            print '| %*s | %*s |' % (-20,project,-20,len([s for s in projects_assessors_list if s['project_id']==project]))
        print ' ---------------------------------------------\n'
    return projects_assessors_list

def filter_assessors(options, obj_list):
    subjects = get_option_list(options.subjects)
    sessions = get_option_list(options.sessions)
    proctypes = get_option_list(options.proctypes)
    #FS for FREESURFER
    if proctypes and proctypes != 'all' and 'FS' in proctypes:
        proctypes.remove('FS')
        proctypes.append('FreeSurfer')
    status = options.formerstatus
    qastatus = options.qastatus
    if subjects and subjects != 'all':
        obj_list = filter(lambda x: x['subject_label'] in subjects, obj_list)
    if sessions and sessions != 'all':
        obj_list = filter(lambda x: x['session_label'] in sessions, obj_list)
    if proctypes and proctypes != 'all':
        obj_list = filter(lambda x: x['proctype'] in proctypes, obj_list)
    if status and status != 'all':
        if qastatus:
            obj_list = filter(lambda x: x['qcstatus'] in status, obj_list)
        else:
            obj_list = filter(lambda x: x['procstatus'] in status, obj_list)  
    return obj_list

def delete_Out_Resource(assessor_obj,resource_label):
    deleted=False
    count=0
    while count<3 and deleted==False:
        try:
            if assessor_obj.out_resource(resource_label).exists():
                assessor_obj.out_resource(resource_label).delete()
            deleted=True
        except Exception as e:
            if isinstance(e,KeyboardInterrupt):
                sys.exit()
            else:
                sys.stdout.write('     ->WARNING: Timing Out while deleting: Resource Too big.Trying again.\n')
                count+=1
    
    if not deleted:
        sys.stdout.write('     ->WARNING: Can not remove resource '+resource_label+'. Deleting file by file.\n')
        try:
            for fname in assessor_obj.out_resource(resource_label).files().get()[:]:
                assessor_obj.out_resource(resource_label).file(fname).delete()
                assessor_obj.out_resource(resource_label).delete()
        except Exception as e:
            if isinstance(e,KeyboardInterrupt):
                sys.exit()
            else:
                sys.stdout.write('     ->ERROR: deleting file by file for the resource '+resource_label+'\n')
                print e
    sys.stdout.write('     ->Resource '+resource_label+' deleted\n')

########################################## SWITCH JOB/QC STATUS FS/DEFAULT PROC ##########################################
def set_qastatus_FS(assessor_obj,status):
    assessor_obj.attrs.set('fs:fsData/validation/status',status)
    sys.stdout.write('   -QA Status on Assessor '+assessor_obj.label()+' changed to '+status+'\n')
    
def set_qastatus_default(assessor_obj,status):
    assessor_obj.attrs.set('proc:genProcData/validation/status',status)
    sys.stdout.write('   -QA Status on Assessor '+assessor_obj.label()+' changed to '+status+'\n')
    
def set_status_FS(assessor_obj,status):
    assessor_obj.attrs.set('fs:fsData/procstatus',status)
    if status=='NEED_TO_RUN' or status=='NEED_INPUTS':
        assessor_obj.attrs.set('fs:fsData/validation/status','Job Pending')
        assessor_obj.attrs.set('fs:fsData/jobid',' ')
        assessor_obj.attrs.set('fs:fsData/memused',' ')
        assessor_obj.attrs.set('fs:fsData/walltimeused',' ')
        assessor_obj.attrs.set('fs:fsData/jobnode',' ')
    #display
    sys.stdout.write('   -JOB Status on Assessor '+assessor_obj.label()+' changed to '+status+'\n')

def set_status_default(assessor_obj,status):
    assessor_obj.attrs.set('proc:genProcData/procstatus',status)
    if status=='NEED_TO_RUN' or status=='NEED_INPUTS':
        assessor_obj.attrs.set('proc:genProcData/validation/status','Job Pending')
        assessor_obj.attrs.set('proc:genProcData/jobid',' ')
        assessor_obj.attrs.set('proc:genProcData/memused',' ')
        assessor_obj.attrs.set('proc:genProcData/walltimeused',' ')
        assessor_obj.attrs.set('proc:genProcData/jobnode',' ')
    #display
    sys.stdout.write('   -JOB Status on Assessor '+assessor_obj.label()+' changed to '+status+'\n')

def set_NEED_INPUTS_proctype(xnat,assessor_dict,needinputs_proctypes):
    list_NI_assessors=XnatUtils.list_assessors(xnat, assessor_dict['project_id'], assessor_dict['subject_label'], assessor_dict['session_label'])
    list_NI_assessors=filter(lambda x: x['proctype'] in needinputs_proctypes, list_NI_assessors)
    for assessor in list_NI_assessors:
        assessor_obj=xnat.select('/project/'+assessor['project_id']+'/subjects/'+assessor['subject_id']+'/experiments/'+assessor['session_id']+'/assessors/'+assessor['label'])
        if assessor_obj.exists():
            if isFreeSurfer(assessor['proctype']): set_status_FS(assessor_obj,'NEED_INPUTS')
            else: set_status_default(assessor_obj,'NEED_INPUTS')
            for resource in XnatUtils.list_assessor_out_resources(xnat,assessor['project_id'],assessor['subject_id'],assessor['session_id'],assessor['label']):
                delete_Out_Resource(assessor_obj,resource['label'])

########################################## SET ASSESSOR ##########################################                 
def set_status_assessor(xnat,assessor_dict,needinputs,status,deleteR,qastatus):
    assessor_obj=xnat.select('/project/'+assessor_dict['project_id']+'/subjects/'+assessor_dict['subject_label']+'/experiments/'+assessor_dict['session_label']+'/assessors/'+assessor_dict['label'])
    if not assessor_obj.exists():
        print 'ERROR: Assessors '+assessor_dict['label']+' does not exist on XNAT.'
    else:
        if qastatus:
            if isFreeSurfer(assessor_dict['proctype']):
                set_qastatus_FS(assessor_obj,status)
            else:
                set_qastatus_default(assessor_obj,status)
        else:
            if isFreeSurfer(assessor_dict['proctype']): 
                set_status_FS(assessor_obj,status)
            else:
                set_status_default(assessor_obj,status)
            if needinputs:
                sys.stdout.write('  +Setting assessors status to NEED_INPUTS that are linked to '+assessor_obj.label()+' : \n')
                set_NEED_INPUTS_proctype(xnat,assessor_dict,needinputs)
            if deleteR:
                for resource in XnatUtils.list_assessor_out_resources(xnat,assessor_dict['project_id'],assessor_dict['subject_label'],assessor_dict['session_label'],assessor_dict['label']):
                    delete_Out_Resource(assessor_obj,resource['label'])

########################################## SWITCH ALL ##########################################
def Switch_project_status(xnat,options):
    #Variables
    status=options.status
    projects_assessors_list=xnat_list_assessors(xnat,options)
    if not projects_assessors_list:
        print 'INFO: No assessors found.'
    else:
        #For each assessor in the list sorted by the label
        print 'INFO: Switching assessors status:'
        for index,assessor_dict in enumerate(sorted(projects_assessors_list, key=lambda k: k['label'])):
            sys.stdout.write('  + Process '+str(index+1)+'/'+str(len(projects_assessors_list))+' : '+assessor_dict['label']+'\n')
            sys.stdout.flush()
            set_status_assessor(xnat,assessor_dict,options.needinputs,status,options.deleteR,options.qastatus)

########################################## CHECK OPTIONS ##########################################
def check_options(options):
    #Checked argument values if not:
    if options.txtfile:
        if not os.path.exists(options.txtfile):
            print "OPTION ERROR: the file "+options.txtfile+" does not exist."
            return False
    else:
        if not options.txtfile and not options.select and not options.printstatus:
            if not options.project :
                print'OPTION ERROR: No project ID given, please give one with -p options. Use -h to check the options.'
                return False
            if not options.proctypes :
                print'OPTION ERROR: No process type given, please give one with -t options. Use -h to check the options.'
                print'E.G: fMRIQA,dtiQA_v2,FreeSurfer'
                return False
    
    if not options.status and not options.printstatus:
        print 'OPTION ERROR: No status given, please give one with -s options. Use -h to check the options.'
        return False
    
    if options.deleteR:
        print'OPTION WARNING: The resources/files on the process will be deleted before changing the status since you used the option -d / --deleteR.'
    
    return True
                            
########################################## MAIN DISPLAY FUNCTION ##########################################
def Main_display(parser):
    (options,_) = parser.parse_args()    
    #Display:
    print '################################################################'
    print '#                   XNATSWITCHPROCESSSTATUS                    #'
    print '#                                                              #'
    print '# Developed by the masiLab Vanderbilt University, TN, USA.     #'
    print '# If issues, email benjamin.c.yvernault@vanderbilt.edu         #'
    print '# Usage:                                                       #'
    print '#     Change assessor job/quality control status               #'
    print '# Parameters :                                                 #'
    if options=={'printstatus':False,'status': None, 'qastatus': False, 'needinputs': None, 'proctypes': None, 'txtfile': None, 'project': None, 'sessions': None, 'deleteR': False, 'formerstatus': None, 'select': None, 'subjects': None}:
        print '#     No Arguments given                                       #'
        print '#     Use "XnatSwitchProcessStatus -h" to see the options      #'
        print '################################################################'
        parser.print_help()
        sys.exit()
    else:        
        if options.txtfile:
            print '#     %*s -> %*s#' %(-20,'File txt',-33,get_proper_str(options.txtfile,True))
        elif options.select:
            print '#     %*s -> %*s#' %(-20,'Selected Process',-33,get_proper_str(options.select,True))
        else:
            if options.project:
                print '#     %*s -> %*s#' %(-20,'Project(s)',-33,get_proper_str(options.project))
            #Subjects
            if options.subjects:
                print '#     %*s -> %*s#' %(-20,'Subject(s)',-33,get_proper_str(options.subjects))
            #Experiment
            if options.sessions:
                print '#     %*s -> %*s#' %(-20,'Session(s)',-33,get_proper_str(options.sessions))
            #Processes    
            if options.proctypes:
                print '#     %*s -> %*s#' %(-20,'Process Types',-33,get_proper_str(options.proctypes))
            if options.needinputs:
                print '#     %*s -> %*s#' %(-20,'NEED_INPUTS Types',-33,get_proper_str(options.needinputs))
        if options.status:
            print '#     %*s -> %*s#' %(-20,'New Status',-33,options.status)
        if options.formerstatus:
            print '#     %*s -> %*s#' %(-20,'Previous Status',-33,get_proper_str(options.formerstatus))
        if options.qastatus:
            print '#     %*s -> %*s#' %(-20,'Change QA status',-33,'on')
        if options.deleteR:
            print '#     %*s -> %*s#' %(-20,'Delete resources',-33,'on')
        if options.printstatus:
            print '#     %*s -> %*s#' %(-20,'Print Status',-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_description():
    usage="""usage: %prog [options]\nWhat is the script doing : 
    *Switch/Set the status for assessors on XNAT specify by their proctype.
    \nExamples:
    *See status managed by DAX: XnatSwitchProcessStatus --printstatus
    *Set all fMRIQA to a specific status Error for a project: XnatSwitchProcessStatus -p PID -s Error -t fMRIQA
    *Set all Multi_Atlas that have the status JOB_FAILED to NEED_TO_RUN to have the processes run again: XnatSwitchProcessStatus -p PID -f JOB_FAILED -t Multi_Atlas -s NEED_TO_RUN
    *Set all VBMQA to NEED_TO_RUN for a project and delete resources: XnatSwitchProcessStatus -p PID -s NEED_TO_RUN -t VBMQA -d
    *Set all VBMQA to NEED_TO_RUN, delete resources, and set linked assessors fMRI_Preprocess to NEED_INPUTS: XnatSwitchProcessStatus -p PID -s NEED_TO_RUN -t VBMQA -d -n fMRI_Preprocess
    *Set all dtiQA_v2 qa status to Passed for a project: XnatSwitchProcessStatus -p PID -s Passed -t dtiQA_v2 --qc
    *Set FreeSurfer for a specific project/subject to NEED_INPUTS: XnatSwitchProcessStatus -p PID --subj 123 -s NEED_INPUTS -t FreeSurfer"""
    return usage

########################################## OPTIONS ##########################################
def parse_args():
    from optparse import OptionParser
    usage =get_description()
    parser = OptionParser(usage=usage)
    parser.add_option("--select", dest="select",default=None,
                    help="Give the process label that you want to change the status. E.G : BLSA-x-BLSA_0000-x-BLSA_0000_00-x-FreeSurfer", metavar="STRING")   
    parser.add_option("-x","--txtfile",dest="txtfile",default=None,
                    help="File txt with at each line the label of the assessor where the status need to be changed. E.G for label: project-x-subject-x-experiment-x-scan-x-process_name.", metavar="STRING")  
    parser.add_option("-p", "--project", dest="project",default=None,
                    help="Project ID on XNAT or list of Project ID", metavar="COMMA_SEPARATED_LIST")
    parser.add_option("--subj", dest="subjects",default=None,
                  help="Change Status for only this subject/list of subjects. E.G: --subj VUSTP2,VUSTP3", metavar="COMMA_SEPARATED_LIST")
    parser.add_option("--sess", dest="sessions",default=None,
                  help="Change Status for only this session/list of sessions. Use the options --subj with it. E.G: --exp VUSTP2a,VUSTP3b", metavar="COMMA_SEPARATED_LIST")
    parser.add_option("-s", "--status", dest="status",default=None,
                    help="Status you want to set on the Processes. E.G: 'NEED_TO_RUN'", metavar="STRING")
    parser.add_option("-f","--formerStatus", dest="formerstatus",default=None,
                    help="Status that need to be changed to the new one. E.G: 'JOB_FAILED'", metavar="STRING")           
    parser.add_option("-t","--type", dest="proctypes",default=None,
                    help="Process type you want the status to changed. E.G: fMRIQA,dtiQA_v2. You can use 'all' for all of them. If not use, will change the status for the processes with the right Type.", metavar="COMMA_SEPARATED_LIST")  
    parser.add_option("-n","--Needinputs", dest="needinputs",default=None,
                    help="Process type that will need inputs if you change the proctype given to the option -t.", metavar="COMMA_SEPARATED_LIST") 
    parser.add_option("-d","--deleteR",dest="deleteR",action="store_true", default=False,
                    help="Delete the resources present on the process.")
    parser.add_option("--qc",dest="qastatus",action="store_true", default=False,
                    help="Change the quality control status on XNAT.")
    parser.add_option("--printstatus",dest="printstatus",action="store_true", default=False,
                    help="Print status used by DAX to manage assessors.")
    return parser

########################################## MAIN FUNCTION ##########################################
if __name__ == '__main__':
    parser=parse_args()
    (options,args) = parser.parse_args()
    #############################
    #Main display:
    Main_display(parser)
    #check options:
    run=check_options(options)
    #############################
    
    #############################
    # RUN                       #
    #############################
    if run:
        #############################
        status=options.status
        needinputs=get_option_list(options.needinputs)
        if needinputs and status!="NEED_TO_RUN":
            print 'Warning: You want to change the status for linked assessors '+options.needinputs+' to NEED_INPUTS but the status set is not NEED_TO_RUN but '+options.status
            print '         The tool will not change the status on the linked assessors.'
        if options.printstatus:
            print "INFO: Status used by DAX package:"
            print " * NEED_INPUTS       - the assessor need inputs to run. Default status when creating the assessor."
            print " * NEED_TO_RUN       - the assessor is ready to run. It will be launch in the next dax_update_open_task."
            print " * JOB_RUNNING       - the assessor is running on the cluster."
            print " * JOB_FAILED        - the assessor failed. You should check the outlog files."
            print " * READY_TO_UPLOAD   - the assessor is waiting in the upload queue to be upload."
            print " * UPLOADING         - the assessor is being uploaded at this instant."
            print " * READY_TO_COMPLETE - the assessor is waiting for dax_update_open_task to set the job information on XNAT (memory/walltime...)"
            print " * COMPLETE          - the assessor is complete. All resources and informations are on XNAT."
            print "QA status: "
            print " * Reproc            - rerun an assessor with new resources (EG: FreeSurfer)."
            print " * passed            - assessor that passed the quality control."
            print " * failed            - assessor that failed the quality control." 
        try:
            xnat = XnatUtils.get_interface()
            if options.select: #SPECIFIC ASSESSOR SELECTED
                print "INFO: Changing Status on assessor "+options.select+" to "+status
                set_status_assessor(xnat,get_assessor_dict(options.select),needinputs,status,options.deleteR,options.qastatus)            
            elif options.txtfile: #TEXTFILE
                assessor_label_list=list()
                input_file = open(options.txtfile, 'r')
                for index,line in enumerate(input_file):
                    assessor_label=line.strip().split('\n')[0]
                    set_status_assessor(xnat,get_assessor_dict(assessor_label),needinputs,status,options.deleteR,options.qastatus)
            elif options.project:  #ALL ASSESSORS FOR PROJECT
                Switch_project_status(xnat,options)
        finally:                
            xnat.disconnect()    
        
    print '===================================================================\n'
