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

'''
Reset sessions to be seen by the nex dax_update.

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

from __future__ import print_function

from builtins import str

import os

from dax import XnatUtils
from dax.errors import XnatToolsUserError
import dax.xnat_tools_utils as utils


__copyright__ = 'Copyright 2013 Vanderbilt University. All Rights Reserved'
__exe__ = os.path.basename(__file__)
__author__ = 'byvernault'
__purpose__ = 'Reset sessions to be seen by the nex dax_update.'
__description__ = """What is the script doing :
   * Reset sessions last update date to update the sessions during
     the next dax_update.

Examples:
   *Reset all sessions:
        Xnatsessionupdate -p PID --all
   *Reset some sessions :
        Xnatsessionupdate -p PID -s 109374,109348
   *Reset for the sessions that have assessors NEED_INPUTS:
        Xnatsessionupdate -p PID -n"""


def get_text_session(xnat, txt_file, project, sessions_list):
    """
    Method to get session list from a textfile after checking their existences

    :param xnat: pyxnat interface
    :param txt_file: text file to export sessions
    :param project: project ID on XNAT
    :param sessions_list: list of sessions dictionaries in the project
    :return: list of sessions labels to edit
    """
    if not os.path.exists(txt_file):
        err = 'text file does not exist: %s'
        raise XnatToolsUserError(__exe__, err % txt_file)

    sessions_labels = list()
    input_file = open(txt_file, 'r')
    for line in input_file:
        label = line.strip().split('\n')[0]
        # check on XNAT:
        exists, session = is_good_session(xnat, project, label, sessions_list)
        if exists:
            sessions_labels.append(session)

    return sessions_labels


def is_good_session(xnat, project, session_label, sessions_list):
    """
    Method to check if a session exists on a project

    :param xnat: pyxnat interface
    :param project: project ID on XNAT
    :param session_label: session label to change
    :param sessions_list: list of sessions dictionaries in the project
    :return: True if subject exists, False otherwise and the session label
    """
    _err = '{label} "{value}" not found on XNAT. Wrong label for {label} in \
text file.'
    labels = session_label.split('-x-')
    if len(labels) > 1:
        if len(labels) != 3:
            err = 'Not enoug information for "%s" in text file. \
The format is project-x-subject-x-session'
            raise XnatToolsUserError(__exe__, err % session_label)
        project = labels[0]
        subject = labels[1]
        session = labels[2]
    else:
        subject = get_subject_id(session_label, sessions_list)
        session = session_label

    project_obj = XnatUtils.select_obj(xnat, project)
    if not project_obj.exists():
        err = _err.format(label='Project', value=project)
        raise XnatToolsUserError(__exe__, err)
    else:
        subject_obj = project_obj.subject(subject)
        if not subject_obj.exists():
            err = _err.format(label='Subject', value=subject)
            raise XnatToolsUserError(__exe__, err)
        else:
            session_obj = subject_obj.experiment(session)
            if session_obj.exists():
                return True, session
            else:
                err = _err.format(label='Session', value=session)
                raise XnatToolsUserError(__exe__, err)


def get_subject_id(session, sessions_list):
    """
    Method to return the subject ID from a list of sessions

    :param session: session label
    :param sessions_list: list of sessions dictionaries for a project on XNAT
    :return: subject id
    """
    subject_id = [x['subject_id'] for x in sessions_list
                  if x['session_label'] == session]
    return subject_id[0]


def change_date(xnat, project, session, sessions_list):
    """
    Method to change the data on all sessions

    :param xnat: pyxnat interface
    :param project: project ID on XNAT
    :param session: session label
    :param sessions_list: list of session dictionaries for the project
    :return: None
    """
    subject = get_subject_id(session, sessions_list)
    session_obj = XnatUtils.select_obj(xnat, project, subject, session)
    # set the date to nothing:
    session_obj.attrs.set('xnat:mrSessionData/original', ' ')
    print('   ->Session changed to need update.')


def change_date_all(xnat, sessions_list):
    """
    Method to change the data on all sessions

    :param xnat: pyxnat interface
    :param sessions_list: list of session dictionaries to modify
    :return: None
    """
    for session_dict in sessions_list:
        session_object = XnatUtils.get_full_object(xnat, session_dict)
        session_object.attrs.set('xnat:mrSessionData/original', ' ')
        msg = '   ->Session "%s" changed to need update.'
        print(msg % (session_dict['label']))


def get_xnat_subject_need_inputs(xnat, project):
    """
    Method to get the sessions labels from XNAT where processes need inputs

    :param xnat: pyxnat interface
    :param project: project ID on XNAT
    :return: unique list of sessions labels on XNAT
    """
    assessors_list = XnatUtils.list_project_assessors(xnat, project)
    assessors_list = [x for x in assessors_list
                      if x['procstatus'] == 'NEED_INPUTS']
    return list(set([x['session_label'] for x in assessors_list]))


def run_xnat_session_update(args):
    """
    Main function for xnat session update.

    :param args: arguments parse by argparse
    """
    too_many_mode = ((args.need_inputs and args.session) or
                     (args.need_inputs and args.txt_file) or
                     (args.txt_file and args.session))
    no_mode = (not args.session and not args.txt_file and
               not args.need_inputs and not args.all)
    if too_many_mode or no_mode:
        err = 'provide one and only one of those options: --session, \
--txtfile, --needinputs, or --all'
        raise XnatToolsUserError(__exe__, err)

    if args.host:
        host = args.host
    else:
        host = os.environ['XNAT_HOST']
    user = args.username

    utils.print_separators()

    with XnatUtils.get_interface(host=host, user=user) as xnat:
        print('INFO: connection to xnat <%s>:' % host)

        projects = args.projects.split(',')
        for ind, project in enumerate(projects):
            p_obj = XnatUtils.select_obj(xnat, project)
            if not p_obj.exists():
                err = "project %s not found on XNAT."
                raise XnatToolsUserError(__exe__, err % project)
            else:
                print('-> Project : %s' % project)
                print('INFO: Getting Sessions from XNAT...')
                sessions = XnatUtils.list_sessions(xnat, project)
                if not sessions:
                    err = "You don't have access to the project: %s."
                    raise XnatToolsUserError(__exe__, err % project)

                if args.all:
                    print('INFO: Setting all XNAT sessions...')
                    change_date_all(xnat, sessions)
                elif args.need_inputs:
                    print('INFO: Setting for session with NEED_INPUTS \
processes...')
                    sessions_to_update = get_xnat_subject_need_inputs(
                        xnat, project)
                elif args.txt_file:
                    msg = ' INFO: Setting XNAT sessions from text file "%s"...'
                    print(msg % args.txt_file)
                    sessions_to_update = get_text_session(
                        xnat, args.txt_file, project, sessions)
                elif args.session:
                    msg = ' INFO: Setting XNAT sessions for %s...'
                    print(msg % args.session)
                    sessions_to_update = args.session.split(',')
                    for session in sessions_to_update:
                        if len([x for x in sessions
                                if x['session_label'] == session]) == 0:
                            err = 'Session "%s" in -s/--session argument not \
found on XNAT.'
                            raise XnatToolsUserError(__exe__, err % session)

                if not args.all:
                    nb_s = str(len(sessions_to_update))
                    print(' INFO: Changing XNAT sessions last update date for \
%s sessions ...' % nb_s)
                    for ind, session in enumerate(sessions_to_update):
                        msg = '  * Session %s/%s : %s'
                        print(msg % (str(ind + 1), nb_s, session))
                        change_date(xnat, project, session, sessions)

    utils.print_end(__exe__)


def add_to_parser(parser):
    """
    Method to add arguments to default parser for xnat_tools in utils.

    :param parser: parser object
    :return: parser object with new arguments
    """
    parser.add_argument("-p", "--project", dest="projects", required=True,
                        help="Projects ID on Xnat.")
    parser.add_argument("-s", "--session", dest="session", default=None,
                        help="Session label on Xnat or list of them.")
    _h = "Change the subject last update date for all the subject with \
processes that have a job status equal to NEED_INPUTS."
    parser.add_argument("-n", "--needinputs", dest="need_inputs",
                        action="store_true", help=_h)
    _h = "File txt with at each line the label of the assessor or just the \
Session label where the Subject date need to be changed. E.G for label: \
project-x-subject-x-experiment-x-scan-x-process_name."
    parser.add_argument("-x", "--txtfile", dest="txt_file", default=None,
                        help=_h)
    parser.add_argument("-a", "--all", dest="all", action="store_true",
                        help="Change for all sessions.")
    return parser


if __name__ == '__main__':
    utils.run_tool(__exe__, __description__, add_to_parser, __purpose__,
                   run_xnat_session_update)
