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

'''
Upload manual edits to TRACULA subject on XNAT - creates an EDITS resource and
uploads control points from dlabel/dmri

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

from __future__ import print_function

import glob
import os
import sys
import time

from dax import XnatUtils, task


__copyright__ = 'Copyright 2013 Vanderbilt University. All Rights Reserved'
__exe__ = os.path.basename(__file__)
__author__ = 'Brian Boyd'
__purpose__ = "Upload TRACULA edits to XNAT"
CPTS_LIST = [
    'fmajor_PP_avg33_mni_bbr_cpts_7.txt',
    'fminor_PP_avg33_mni_bbr_cpts_5.txt',
    'lh.atr_PP_avg33_mni_bbr_cpts_5.txt',
    'lh.cab_PP_avg33_mni_bbr_cpts_4.txt',
    'lh.ccg_PP_avg33_mni_bbr_cpts_5.txt',
    'lh.cst_AS_avg33_mni_bbr_cpts_6.txt',
    'lh.ilf_AS_avg33_mni_bbr_cpts_5.txt',
    'lh.slfp_PP_avg33_mni_bbr_cpts_5.txt',
    'lh.slft_PP_avg33_mni_bbr_cpts_5.txt',
    'lh.unc_AS_avg33_mni_bbr_cpts_5.txt',
    'rh.atr_PP_avg33_mni_bbr_cpts_5.txt',
    'rh.cab_PP_avg33_mni_bbr_cpts_4.txt',
    'rh.ccg_PP_avg33_mni_bbr_cpts_5.txt',
    'rh.cst_AS_avg33_mni_bbr_cpts_6.txt',
    'rh.ilf_AS_avg33_mni_bbr_cpts_5.txt',
    'rh.slfp_PP_avg33_mni_bbr_cpts_5.txt',
    'rh.slft_PP_avg33_mni_bbr_cpts_5.txt',
    'rh.unc_AS_avg33_mni_bbr_cpts_5.txt'
]


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', help='Subjects Directory',
                    default='/tmp')

    return ap.parse_args()


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

    local_subj_path = glob.glob(os.path.join(
        subjects_dir, '*{}*{}*TRACULA_v*'.format(proj_label, sess_label),
        'TRACULA', sess_label))[0]
    if not os.path.exists(local_subj_path):
        print('ERROR:cannot upload, subject not found in local 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))

        # Find the session
        trac_sess = None
        sess_list = XnatUtils.list_sessions(xnat, projectid=proj_label)
        for sess in sess_list:
            if sess['label'] == sess_label:
                trac_sess = sess
                break

        if not trac_sess:
            print('ERROR:session not found on project:%s' % (sess_label))
            sys.exit(1)

        # Find the TRACULA assessor
        assr_list = XnatUtils.list_assessors(
            xnat, proj_label, trac_sess['subject_ID'], trac_sess['ID'])
        trac_list = [x for x in assr_list if x['proctype'] == 'TRACULA_v1']
        if not trac_list:
            print('ERROR:TRACULA not found for project=%s, session=%s'
                  % (proj_label, sess_label))
            sys.exit(1)

        if len(trac_list) == 1:
            trac = trac_list[0]
        elif not proc_suffix:
            print('ERROR:multiple TRACULA runs found, please specify a \
suffix, Project=%s, Session=%s' % (proj_label, sess_label))
            sys.exit(1)
        else:
            # Get filtered list
            trac_list2 = [x for x in trac_list
                          if x['assessor_label'].endswith(proc_suffix)]
            if not trac_list2:
                print('ERROR:TRACULA not found for project=%s, session=%s'
                      % (proj_label, sess_label))
                sys.exit(1)

            if len(trac_list2) == 1:
                trac = trac_list2[0]
            else:
                print('ERROR:multiple TRACULA runs found with specified \
suffix, Project=%s, Session=%s' % (proj_label, sess_label))
                sys.exit(1)

        # Upload edits
        trac_assr = XnatUtils.get_full_object(xnat, trac)
        resource = trac_assr.out_resource('EDITS')
        curtime = time.strftime("%Y%m%d-%H%M%S")
        local_dlabel_path = os.path.join(local_subj_path, 'dlabel', 'diff')

        for cpts in CPTS_LIST:
            cpts_file = cpts + '.manual'
            cpts_path = os.path.join(local_dlabel_path, cpts_file)

            if os.path.isfile(cpts_path):
                print('Uploading:%s' % (cpts))
                resource.file(cpts + '.' + curtime).put(cpts_path)
            else:
                print('Did not find edited:%s' % (cpts))

        # Set QC Status to trigger reprocessing
        print('Setting qcstatus to trigger reprocessing...')
        trac_assr.attrs.set(
            XnatUtils.DEFAULT_DATATYPE + '/validation/status', task.REPROC)

        # Done
        print('DONE')
