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

'''
Query through Xnat

@author: Benjamin Yvernault, Electrical Engineering, Vanderbilt University
'''

import os
import sys
from dax import XnatUtils
from datetime import datetime

######################################################################################################
########################################## USEFUL FUNCTIONS ##########################################
######################################################################################################
def display_item(project,subject=None,session=None):
    print'Project: '+project
    if subject:
        print '  +Subject: '+subject
        if session:
            print '    *Session: '+session
    
########################################################################################################
########################################## SPECIFIC FUNCTIONS ##########################################
########################################################################################################
def query_project(xnat,project,all=False):
    for subject in XnatUtils.list_subjects(xnat,project):
        print '  +Subject: '+subject['label']
        if all:
            query_subject(xnat,project,subject['label'],all)
    
def query_subject(xnat,project,subject,all=False):
    for session in XnatUtils.list_experiments(xnat,project,subject):
        print '    *Session: '+session['label']
        if all:
            query_session(xnat,project,subject,session['label'])
            
def query_session(xnat,project,subject,session,all=False):
    print '      ***SCANS***' 
    for scan in XnatUtils.list_scans(xnat,project,subject,session):
        query_scan(xnat,project,subject,session,scan['ID'],scan['type'],all)        
    print '      ***PROCESSES***'
    for assessor in XnatUtils.list_assessors(xnat,project,subject,session):
        query_assessor(xnat,assessor['label'],all)
            
def query_scan(xnat,project,subject,session,scan,scantype,all=False):
    print '       -'+scan+' -- '+scantype
    if all:
        for resource in XnatUtils.list_resources(xnat,project,subject,session,scan):
            print '         ->'+resource['label']

def query_assessor(xnat,assessor_label,all=False):
    labels=assessor_label.split('-x-')
    print '       -'+assessor_label
    if all:
        for out_resource in XnatUtils.list_out_resources(xnat,labels[0],labels[1],labels[2],assessor_label):
            print '         ->'+out_resource['label']

########################################## CHECK OPTIONS ##########################################
def check_options(options):
    # The options :
    if not options.me:
        if not options.project and not options.all and not options.assessor:
            print 'OPTION ERROR: You need to specify a project with the option -p/--project.'
            return False
   
    try:
        xnat = XnatUtils.get_interface()
        #ASSESSOR
        if options.assessor:
            if len(options.assessor.split(','))>1:
                print 'OPTION ERROR: Too much processes labels given to the option.'
                return False
            labels=options.assessor.split('-x-')
            if not len(labels)>1:
                print 'OPTION ERROR: You used the option -a/--assessor with an invalid process label.'
                return False
            else:
                if not xnat.select('/project/'+labels[0]+'/subjects/'+labels[1]+'/experiments/'+labels[2]+'/assessors/'+options.assessor).exists():
                    print 'OPTION ERROR: You used the option -a/--assessor with an not existing process label.'
                    return False
        #PROJECT
        if options.project:
            if len(options.project.split(','))>1:
                print 'OPTION ERROR: Too much projects IDs given to the option.'
                return False
            P=xnat.select('/project/'+options.project)
            if not P.exists():
                print 'OPTION ERROR: You used the option -p/--project with an not existing project ID.'
                return False
            #SUBJECT
            if options.subject:
                if len(options.subject.split(','))>1:
                    print 'OPTION ERROR: Too much subjects labels given to the option.'
                    return False
                S=P.subject(options.subject)
                if not S.exists():
                    print 'OPTION ERROR: You used the option -s/--subject with an not existing subject label.'
                    return False
                #EXPERIMENT
                if options.experiment:
                    if len(options.experiment.split(','))>1:
                        print 'OPTION ERROR: Too much experiment labels given to the option.'
                        return False
                    E=S.experiment(options.experiment)
                    if not E.exists():
                        print 'OPTION ERROR: You used the option -e/--experiment with an not existing experiment label.'
                        return False
                    #SCAN
                    if options.scan:
                        if len(options.scan.split(','))>1:
                            print 'OPTION ERROR: Too much scans IDs given to the option.'
                            return False
                        C=E.scan(options.scan) 
                        if not C.exists():
                            print 'OPTION ERROR: You used the option -c/--scan with an not existing scan ID.'
                            return False
    finally:                                        
        xnat.disconnect()
                
    return True

########################################## MAIN DISPLAY FUNCTION ##########################################   
def Main_display(parser):
    (options,args)=parser.parse_args()
    print '################################################################'
    print '#                          XNATQUERY                           #'
    print '#                                                              #'
    print '# Developed by the masiLab Vanderbilt University, TN, USA.     #'
    print '# If issues, email benjamin.c.yvernault@vanderbilt.edu         #'
    print '# Usage:                                                       #'
    print '#     Query through XNAT at the level you want                 #'
    print '# Parameters :                                                 #'
    if options=={'me': False, 'all': False, 'scan': None, 'project': None, 'experiment': None, 'assessor': None, 'subject': None}:
        print '#     No Arguments given                                       #'
        print '#     See the help bellow or Use "Xnatquery" -h                #'
        print '################################################################'
        parser.print_help()
        sys.exit()
    else:
        if options.me:
            print '#     %*s#' %(-57,'->List of projects on Xnat you have access to.')
        else:
            if options.assessor:
                print '#     %*s -> %*s#' %(-20,'Query level',-33,'resources')
                print '#     %*s -> %*s#' %(-20,'Process',-33,options.assessor)
            elif not options.project:
                    print '#     %*s -> %*s#' %(-20,'Query level',-33,'projects')
            else:
                if not options.subject:
                    print '#     %*s -> %*s#' %(-20,'Query level',-33,'subjects')
                    print '#     %*s -> %*s#' %(-20,'Project',-33,options.project)
                else:
                    if not options.experiment:
                        print '#     %*s -> %*s#' %(-20,'Query level',-33,'sessions')
                        print '#     %*s -> %*s#' %(-20,'Project',-33,options.project)
                        print '#     %*s -> %*s#' %(-20,'Subject',-33,options.subject)
                    else:
                        if not options.scan or not options.assessor:
                            print '#     %*s -> %*s#' %(-20,'Query level',-33,'scans/processes')
                            print '#     %*s -> %*s#' %(-20,'Project',-33,options.project)
                            print '#     %*s -> %*s#' %(-20,'Subject',-33,options.subject)
                            print '#     %*s -> %*s#' %(-20,'Session',-33,options.experiment)
                        elif options.scan:
                            print '#     %*s -> %*s#' %(-20,'Query level',-33,'resources')
                            print '#     %*s -> %*s#' %(-20,'Project',-33,options.project)
                            print '#     %*s -> %*s#' %(-20,'Subject',-33,options.subject)
                            print '#     %*s -> %*s#' %(-20,'Session',-33,options.experiment)
                            print '#     %*s -> %*s#' %(-20,'Scan',-33,options.scan)
                            
                if options.all:
                    print '#     %*s -> %*s#' %(-20,'All',-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+="   *Query on Xnat at any level.\n"
    #Example
    usage+="Examples:\n"
    usage+="   *Show all the projects you have access to: Xnatquery --me\n"
    usage+="   *Show all projects: Xnatquery --all\n"
    usage+="   *Query a specific level (example scan/assessors for a session): Xnatquery -p PID -s 109873 -e 109873\n"
    usage+="   *Query a specific level with all objects under it : Xnatquery -p PID -s 109873 --all\n"
    return usage
    
def parse_args():
    from optparse import OptionParser
    usage = get_usage()
    parser = OptionParser(usage=usage)
    parser.add_option("-p", "--project", dest="project",default=None,
                  help="project ID on Xnat or 'all' to see all the project.", metavar="PROJECT_ID")
    parser.add_option("-s", "--subject", dest="subject",default=None,
                  help="Subject label on Xnat", metavar="SUBJECT_LABEL")
    parser.add_option("-e", "--experiment", dest="experiment",default=None,
                  help="Session label on Xnat", metavar="EXPERIMENT_LABEL")
    parser.add_option("-a", "--assessor", dest="assessor",default=None,
                  help="Assessor/Process label on XNAT. E.G: VUSTP-x-VUSTP1-x-VUSTP1a-x-FS", metavar="ASSESSOR_LABEL")
    parser.add_option("-c", "--scan", dest="scan",default=None,
                  help="Scan ID on Xnat.", metavar="SCAN_ID")
    
    parser.add_option("--all", dest="all",action="store_true", default=False,
                  help="Print all the objects on XNAT from the level you are at.", metavar="")
    
    parser.add_option("--me", dest="me",action="store_true", default=False,
                  help="Give the projects ID that you have access.", metavar="")
    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:
        #############################
        #Extract values from Options
        project = options.project
        subject = options.subject
        session = options.experiment
        scan = options.scan
        assessor = options.assessor
        
        """ Query on Xnat to give the project / subject / experiment / scan / resources """
        # Connection to Xnat
        try:
            xnat = XnatUtils.get_interface()
            
            #give the name of the project where you have access
            if options.me:
                print 'List of projects on XNAT you have access to:'
                print '---------------------------------------'
                post_uri_project = '/REST/projects/'
                for project in xnat._get_json(post_uri_project):
                    if XnatUtils.list_subjects(xnat,project['ID']):
                        print '%*s : %*s' % (-20, project['ID'], -30, project['name'])
                print '---------------------------------------'
            #level of query
            else:
                #check assessor and scan, if nothing choose, set them to True
                if not project and options.all and not assessor:
                    project='all'
                
                #if all projects
                if assessor:
                    labels=assessor.split('-x-')
                    display_item(labels[0],labels[1],labels[2])
                    query_assessor(xnat,assessor,True)
                    
                elif project=='all':
                    post_uri_project = '/REST/projects'
                    project_list = xnat._get_json(post_uri_project)
                    for project in project_list:
                        print '%*s : %*s' % (-20, project['ID'], -30, project['name'])
                        if not XnatUtils.list_subjects(xnat,project['ID']):
                            print "  ->You don't have access to this project."
                else:
                    if subject:
                        if session:
                            display_item(project,subject,session)
                            if scan:
                                scantype=xnat.select('/project/'+project+'/subjects/'+subject+'/experiments/'+session+'/scans/'+scan).attrs.get('series_description')
                                query_scan(xnat,project,subject,session,scan,scantype,True)
                            else:
                                query_session(xnat,project,subject,session,options.all)
                        else:
                            display_item(project,subject)
                            query_subject(xnat,project,subject,options.all)
                    else:
                        display_item(project)
                        query_project(xnat,project,options.all)
            
        finally:                                        
            xnat.disconnect()
    
    print '===================================================================\n'
