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

'''
Download a FreeSurfer subject from XNAT

@author: Brian D. Boyd, Psychiatry, Vanderbilt University
'''

import os
import sys
import shutil

from pyxnat import Interface

from dax import XnatUtils

def parse_args():
    from argparse import ArgumentParser
    ap = ArgumentParser(prog='fsdownload', description="Download FreeSurfer subject from XNAT")
    ap.add_argument('project', help='Project Label')
    ap.add_argument('session', help='Session Label')
    return ap.parse_args()
    
if __name__ == '__main__':    
    args = parse_args()
    proj_label = args.project
    sess_label = args.session
    fs = None
                      
    try:
        # Environs
        xnat_user = os.environ['XNAT_USER']
        xnat_pass = os.environ['XNAT_PASS']
        xnat_host = os.environ['XNAT_HOST']
        subjects_dir = os.environ['SUBJECTS_DIR']

    except KeyError as e:
        print "You must set the environment variable %s" % str(e)
        sys.exit(1)  
        
    if os.path.exists(subjects_dir+'/'+sess_label):
        print 'ERROR:cannot download, session already exists in FreeSurfer subjects directory.'
        sys.exit(1)
        
    xnat = Interface(xnat_host, xnat_user, xnat_pass)
    
    #TODO: check that project exists 
    
    # Find the FreeSurfer assessor
    sess_list = XnatUtils.list_experiments(xnat, projectid=proj_label)
    for sess in sess_list:
        if sess['label'] == sess_label:
            assr_list = XnatUtils.list_assessors(xnat, proj_label, sess['subject_ID'], sess['ID'])
            for assr in assr_list:
                # Skip if not FreeSurfer
                if assr['proctype'] == 'FreeSurfer': fs = assr
                    
    if fs == None:
        print('ERROR:FreeSurfer not found for project='+proj_label+ ', session='+sess_label)
        sys.exit(1)
                        
    # Download it
    assr_label = fs['assessor_label']
    out_zip = subjects_dir +'/'+sess_label+'.zip'
    if os.path.exists(subjects_dir+'/'+assr_label): 
        print 'ERROR:cannot download, fs already exists:'+subjects_dir+'/'+assr_label
        sys.exit(1)        
        
    if os.path.exists(out_zip):
        print 'ERROR:cannot download, zip already exists:'+out_zip
        sys.exit(1)
        
    print 'Downloading:'+assr_label
    cmd = 'curl -qu '+xnat_user+':'+xnat_pass+' '+xnat_host
    cmd += '/data/archive/projects/'+fs['project_id']+'/subjects/'+fs['subject_id']
    cmd += '/experiments/'+fs['session_id']+'/assessors/'+fs['assessor_id']
    cmd += '/out/resources/DATA/files?format=zip > '+out_zip
    os.system(cmd)
    
    if os.path.exists(subjects_dir+'/'+sess_label):
        print 'ERROR:cannot unzip, already exists:'+subjects_dir+'/'+sess_label
        sys.exit(1)
        
    # Unzip 
    os.chdir(subjects_dir)
    cmd = 'unzip -q '+out_zip
    os.system(cmd)
 
    print 'Moving files...'
    src_assr_label = subjects_dir+'/'+assr_label+'/out/resources/DATA/files/'+assr_label
    src_sess_label = subjects_dir+'/'+sess_label+'/assessors/'+assr_label+\
        '/out/resources/DATA/files/'+assr_label
    src_sess_id = subjects_dir+'/'+sess_label+'/assessors/'+fs['assessor_id']+\
        '/out/resources/DATA/files/'+assr_label
        
    # Find the FS subject directory in the downloaded directory
    if os.path.exists(src_assr_label):
        src  = src_assr_label
    elif os.path.exists(src_sess_label): 
        src = src_sess_label
    elif os.path.exists(src_sess_id): 
        src = src_sess_id     
    else: 
        print('ERROR:downloaded subject does not exist')
        sys,exit()

    # Move the subdir containing FS subject up to level of SUBJECTS_DIR, renaming to session label
    dest = subjects_dir+'/'+sess_label
    shutil.move(src,dest)

    # Delete the downloaded directory
    if os.path.exists(subjects_dir+'/'+assr_label+'/'):
        shutil.rmtree(subjects_dir+'/'+assr_label+'/')    
    elif os.path.exists(subjects_dir+'/'+sess_label+'/assessors/'): 
        shutil.rmtree(subjects_dir+'/'+sess_label+'/assessors/')    
        
    # Delete the downloaded zip
    os.remove(out_zip)
