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

'''
Display information on a project from Xnat

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

from __future__ import print_function

from builtins import str

from datetime import datetime
import os

from dax import XnatUtils, task
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__ = 'Display information on a XNAT project.'
__description__ = '''What is the script doing :
   * Generate a report for a XNAT project displaying scans/assessors
     information.

Examples:
    * See the information for project TEST:
        Xnatinfo TEST
'''

STATUS_DICT = {
    task.NEED_INPUTS: 1,
    task.NEED_TO_RUN: 2,
    task.JOB_RUNNING: 3,
    task.JOB_FAILED: 4,
    task.READY_TO_UPLOAD: 5,
    task.UPLOADING: 6,
    task.READY_TO_COMPLETE: 7,
    task.COMPLETE: 8,
    task.NO_DATA: 9,
    'UNKNOWN': 10
}

INFO_TEMPLATE = '''
Information for project {project} on Xnat :
Date: {date}
==========================================================================

Project Info:
--------------------------------------------------------------------------
Count:
---------------------------------------
Subjects            : {nb_s}
Sessions            : {nb_se}
Scans               : {nb_sc}
Assessors/Processes : {nb_a}
---------------------------------------

{scan_str}

Process info :
--------------------------------------------------------------------------
  Process type              | Count |   C   |   RC  |   U   |   RU  |\
   JF  |   JR  |  NTR  |   NI  |   ND  |   UN
  --------------------------------------------------------------------------\
---------------------------------------
{assessor_info}
--------------------------------------------------------------------------
Status:
  C = COMPLETE
  RC = READY_TO_COMPLETE
  U = UPLOADING
  RU = READT_TO_UPLOAD
  JF = JOB_FAILED
  JR = JOB_RUNNING
  NTR = NEED_TO_RUN
  NI = NEED_INPUTS
  ND = NO_DATA
  UN = UNKNOWN

{extra}'''

SCAN_TEMPLATE = '''Scan info :
--------------------------------------------------------------------------
  Scan type                                                    | Count
  ---------------------------------------------------------------------
{scan_info}
  ---------------------------------------------------------------------

  List of unusable scan :
  -----------------------
{scan_unusable}

--------------------------------------------------------------------------
'''

EXTRA_TEMPLATE = '''  List of UNKNOWN status:
  -----------------------
  {unknown}
  -----------------------
  {run_fail}
--------------------------------------------------------------------------
'''


def add_process_to_dict(dictionary, proctype, procstatus):
    """
    Method to add process to the dictionary for report

    :param dictionary: dictionary to add process to
    :param proctype: type of process to add
    :param procstatus: status of the process to add
    :return: dictionary updated
    """
    # proctype:
    if proctype in dictionary:
        dictionary[proctype][0] += 1
    else:
        dictionary[proctype] = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

    # procstatus:
    dictionary[proctype][STATUS_DICT[procstatus]] += 1

    return dictionary


def report_project(xnat, project, ignore_scans, ignore_unusable, running,
                   failed, output_file):
    """
    Main Method to display the report for a project

    :param xnat: pyxnat interface
    :param project: project id on xnat
    :param ignorescans: don't display the scan info
    :param ignoreunusable: don't print the unusable
    :param running: print running jobs
    :param failed: print failed jobs
    :param output_file: file to print the report.
    :return: None
    """
    # scan
    scans_unusable = list()
    scans_found = dict()
    unknown_status = list()

    # for all process except FS
    assessors_found = dict()

    # Check project:
    proj_obj = xnat.select('/project/%s' % project)
    if not proj_obj.exists():
        err = 'project %s not found on XNAT'
        raise XnatToolsUserError(__exe__, err % project)

    print('INFO: query through project %s ...' % project)
    # scan loop
    scans_list = XnatUtils.list_project_scans(xnat, project)
    subj_number = len(set([d['subject_label'] for d in scans_list]))
    sess_number = len(set([d['session_label'] for d in scans_list]))
    scan_number = len(scans_list)
    for scan_dict in scans_list:
        if scan_dict['quality'] == 'unusable':
            scans_unusable.append([scan_dict['subject_label'],
                                   scan_dict['session_label'],
                                   scan_dict['ID'],
                                   scan_dict['type']])
        # add the count for the scan type:
        if scan_dict['type'] in scans_found:
            scans_found[scan_dict['type']] += 1
        else:
            scans_found[scan_dict['type']] = 1

    # assessor loop
    assessors_list = XnatUtils.list_project_assessors(xnat, project)
    assessor_number = len(assessors_list)
    for assessor_dict in assessors_list:
        # add to dictionary of process
        if assessor_dict['procstatus'] in list(STATUS_DICT.keys()):
            assessors_found = add_process_to_dict(assessors_found,
                                                  assessor_dict['proctype'],
                                                  assessor_dict['procstatus'])
        else:
            unknown_status.append(assessor_dict['procstatus'])
            assessors_found = add_process_to_dict(assessors_found,
                                                  assessor_dict['proctype'],
                                                  'UNKNOWN')

    # Scans display
    scans_str = 'Scans info: skipped...'
    if not ignore_scans:
        scan_info = list()
        for key in sorted(scans_found):
            _format = '  %*s | %*s'
            scan_info.append(_format % (-60, key[:60], -30, scans_found[key]))
        unusables = ['  skipped...']
        if not ignore_unusable:
            _format = '  %*s | %*s | %*s | %*s'
            unusables = [_format % (-20, 'Subject', -20, 'Experiment',
                                    -20, 'Scan', -20, 'Type')]
            for _sc in sorted(scans_unusable):
                unusables.append(_format % (-20, _sc[0], -20, _sc[1],
                                            -20, _sc[2], -20, _sc[3]))

            if len(unusables) == 1:
                unusables.append('  None')

        scans_str = SCAN_TEMPLATE.format(scan_info='\n'.join(scan_info),
                                         scan_unusable='\n'.join(unusables))

    # Assessors display
    as_len = 25
    _format = '  %*s | %*s | %*s | %*s | %*s | %*s | %*s | %*s | %*s | %*s | \
%*s | %*s '
    assrs_info = []
    for assessor_dict in assessors_list:
        if len(assessor_dict['proctype']) > as_len:
            as_len = len(assessor_dict['proctype'])
    for key in sorted(assessors_found):
        assrs_info.append(
            _format %
            (-1 * as_len, key,
             -5, assessors_found[key][0],
             -5, assessors_found[key][STATUS_DICT['COMPLETE']],
             -5, assessors_found[key][STATUS_DICT['READY_TO_COMPLETE']],
             -5, assessors_found[key][STATUS_DICT['UPLOADING']],
             -5, assessors_found[key][STATUS_DICT['READY_TO_UPLOAD']],
             -5, assessors_found[key][STATUS_DICT['JOB_FAILED']],
             -5, assessors_found[key][STATUS_DICT['JOB_RUNNING']],
             -5, assessors_found[key][STATUS_DICT['NEED_TO_RUN']],
             -5, assessors_found[key][STATUS_DICT['NEED_INPUTS']],
             -5, assessors_found[key][STATUS_DICT['NO_DATA']],
             -5, assessors_found[key][STATUS_DICT['UNKNOWN']]))

    extras = list()
    if unknown_status:
        unknown_status = list(set(unknown_status))
        for status in sorted(unknown_status):
            extras.append('  %s' % (status))
    if not extras:
        extras = ['None']
    run_fail = list()
    if running:
        run_fail.append('\nRUNNING JOBS:')
        has_running = False
        for assessor_dict in assessors_list:
            if assessor_dict['procstatus'] == task.JOB_RUNNING:
                has_running = True
                run_fail.append(assessor_dict['label'])
        if not has_running:
            run_fail.append('None')

    if failed:
        run_fail.append('FAILED JOBS:')
        has_failed = False
        for assessor_dict in assessors_list:
            if assessor_dict['procstatus'] == task.JOB_FAILED:
                has_failed = True
                run_fail.append(assessor_dict['label'])
        if not has_failed:
            run_fail.append('None')
    extra_str = EXTRA_TEMPLATE.format(unknown='\n'.join(extras),
                                      run_fail='\n'.join(run_fail))

    # Report:
    report_str = INFO_TEMPLATE.format(
        project=project,
        date=str(datetime.now()),
        nb_s=subj_number,
        nb_se=sess_number,
        nb_sc=scan_number,
        nb_a=assessor_number,
        scan_str=scans_str,
        assessor_info='\n'.join(assrs_info),
        extra=extra_str,
    )

    # Print or write in files:
    if output_file:
        folder = os.path.dirname(os.path.abspath(output_file))
        if not os.path.exists(folder):
            warn = 'Warning: the path selected %s does not exist for the file \
text.\n'
            print(warn % (folder))
            print(report_str)
        else:
            file_txt = os.path.abspath(output_file)
            print('INFO: Writing the report in the file: %s' % (file_txt))
            with open(file_txt, 'w') as f_write:
                f_write.write(report_str)
    else:
        print(report_str)


def run_xnat_info(args):
    """
    Main function for xnat check.xnat

    :param args: arguments parse by argparse
    """
    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)
        report_project(xnat, args.project, args.ignore_scans,
                       args.ignore_unusable, args.running,
                       args.failed, args.output_file)

    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(dest='project', help='Project ID on XNAT')
    parser.add_argument("-x", "--filetxt", dest='output_file', default=None,
                        help='Path to a txt file to save the report')
    parser.add_argument('-f', '--failed', dest='failed', action='store_true',
                        help='Add this flag to print out failed jobs',
                        default=False)
    parser.add_argument('-r', '--running', dest='running', action='store_true',
                        help='Add this flag to print out running jobs',
                        default=False)
    parser.add_argument('--ignoreUnusable', dest='ignore_unusable',
                        help='Ignore print statement of unusable scans',
                        action='store_true')
    parser.add_argument('--ignoreScans', dest='ignore_scans',
                        action='store_true',
                        help='Ignore print statement of scans')
    return parser


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