#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''
Download a FreeSurfer subject from XNAT

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

from __future__ import print_function

import os
import sys
import shutil

from dax import XnatUtils


__copyright__ = 'Copyright 2013 Vanderbilt University. All Rights Reserved'
__exe__ = os.path.basename(__file__)
__author__ = 'Brian Boyd'
__purpose__ = "Download FreeSurfer subject from XNAT"
CURL_CMD = '''curl -qu {xuser}:{xpass} {xhost}/data/archive/projects/{proj}/\
subjects/{subj}/experiments/{sess}/assessors/{asse}/out/resources/DATA/\
files?format=zip > {zip}'''


def parse_args():
    """
    Method to parse arguments base on ArgumentParser

    :return: parser object parsed
    """
    from argparse import ArgumentParser
    ap = ArgumentParser(prog=__exe__, description=__purpose__)
    ap.add_argument('--host', dest='host', default=None,
                    help='Host for XNAT. Default: env XNAT_HOST.')
    ap.add_argument('-u', '--username', dest='username', default=None,
                    help='Username for XNAT.')
    ap.add_argument('project', help='Project Label')
    ap.add_argument('session', help='Session Label')
    ap.add_argument('proc_suffix', help='Proc name suffix', nargs='?',
                    default='')
    ap.add_argument('-sd', '--subjects_dir', dest='subjects_dir',
                    help='Subjects Directory',
                    default=os.environ.get('SUBJECTS_DIR', '/tmp'))
    return ap.parse_args()


if __name__ == '__main__':
    args = parse_args()
    proj_label = args.project
    sess_label = args.session
    subjects_dir = args.subjects_dir
    fs = None
    if not args.proc_suffix:
        proc_suffix = ''
    else:
        proc_suffix = args.proc_suffix

    if os.path.exists(os.path.join(subjects_dir, sess_label)):
        print('ERROR:cannot download, session already exists in FreeSurfer \
subjects directory.')
        sys.exit(1)

    if args.host:
        host = args.host
    else:
        host = os.environ['XNAT_HOST']
    user = args.username
    with XnatUtils.get_interface(host=host, user=user) as xnat:
        print('INFO: connection to xnat <%s>:' % (host))

        # TODO: check that project exists

        # Find the FreeSurfer assessor
        sess_list = XnatUtils.list_experiments(xnat, projectid=proj_label)
        fs_list = []
        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_list.append(assr)

        if not fs_list:
            print('ERROR:FreeSurfer not found for project=%s, session=%s'
                  % (proj_label, sess_label))
            sys.exit(1)

        if len(fs_list) == 1:
            fs = fs_list[0]
        else:
            if not proc_suffix:
                print('ERROR:mutliple FreeSurfer runs found, you must specify an \
    ID suffix, project={}, session={}'.format(proj_label, sess_label))
                sys.exit(1)

            fs_list2 = []
            for assr in fs_list:
                if assr['assessor_label'].endswith(proc_suffix):
                    fs_list2.append(assr)

                    if not fs_list2:
                        print('ERROR:FreeSurfer not found for project=%s, \
    session=%s' % (proj_label, sess_label))
                        sys.exit(1)

                    if len(fs_list2) == 1:
                        fs = fs_list2[0]
                    else:
                        print('ERROR:mutliple FreeSurfer runs found with \
    specified suffix, project=%s, session=%s' % (proj_label, sess_label))
                        sys.exit(1)

        # Download it
        assr_label = fs['assessor_label']
        out_zip = os.path.join(subjects_dir, sess_label + '.zip')
        if os.path.exists(os.path.join(subjects_dir, assr_label)):
            print('ERROR:cannot download, fs already exists: %s'
                  % (os.path.join(subjects_dir, assr_label)))
            sys.exit(1)

        if os.path.exists(out_zip):
            print('ERROR:cannot download, zip already exists: %s' % (out_zip))
            sys.exit(1)

        print('Downloading:{}'.format(assr_label))
        cmd = CURL_CMD.format(xuser=xnat.user, xpass=xnat.pwd, xhost=xnat.host,
                              proj=fs['project_id'], subj=fs['subject_id'],
                              sess=fs['session_id'], asse=fs['assessor_id'],
                              zip=out_zip)
        os.system(cmd)

        if os.path.exists(os.path.join(subjects_dir, sess_label)):
            print('ERROR:cannot unzip, already exists: %s'
                  % (os.path.join(subjects_dir, sess_label)))
            sys.exit(1)

        # Unzip
        os.chdir(subjects_dir)
        cmd = 'unzip -q %s' % (out_zip)
        os.system(cmd)

        # Determine format of unzipped data
        data_dir = os.path.join(subjects_dir, assr_label, 'out', 'resources',
                                'DATA', 'files')
        if os.path.exists(os.path.join(data_dir, assr_label)):
            # <assr_label>/out/resources/DATA/files/<assr_label>
            src = os.path.join(data_dir, assr_label)
        elif os.path.exists(os.path.join(data_dir, 'Subjects', assr_label)):
            # <assr_label>/out/resources/DATA/files/Subjects/assr_label
            src = os.path.join(data_dir, 'Subjects', assr_label)
        elif os.path.exists(os.path.join(data_dir, 'mri')):
            # <assr_label>/out/resources/DATA/files
            src = data_dir
        else:
            print('ERROR:failed to find FreeSurfer data in downloaded files.')
            sys.exit()

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

        # Delete the downloaded directory
        if os.path.exists(os.path.join(subjects_dir, assr_label)):
            shutil.rmtree(os.path.join(subjects_dir, assr_label))

        # Delete the downloaded zip
        os.remove(out_zip)
