#!/usr/bin/env python
"""
     _  _  _ .  _  |_     . |_  | a platform for bash pipeline development
    _) (_ |  | |_) |_ |_| | |_  | github.com/josephdviviano/scriptuit
               |                |

    scriptuit list                 -- lists all available modules
    scriptuit help module          -- prints help for module
    scriptuit check setup          -- checks installation
    scriptuit check inputs         -- checks input files for experiment
    scriptuit generate             -- generate master script
    scriptuit render script output -- render master script to output script
    scriptuit clean                -- delete output files
"""
import os, sys, stat
from copy import copy
import scriptuit as sit
from fnmatch import fnmatch
import textwrap

DIR_DATA = os.getenv('SCRIPTUIT_DATA')
DIR_MODULES = os.getenv('SCRIPTUIT_MODULES')
CLEARSCREEN = chr(27) + "[2J"

def get_modules(interactive=False, used=None):
    """
    Prints the available modules. If provided with a list of 'used' modules,
    these are stripped from the output.
    """
    if not DIR_MODULES:
        sys.exit('ERROR: Environment variable SCRIPTUIT_MODULES not defined.')

    module_list = [f for f in os.listdir(DIR_MODULES)
                           if os.path.isfile(os.path.join(DIR_MODULES, f))]
    module_list.sort()

    # remove used modules from list
    if used:
        for u in used:
            try:
                module_list.remove(u)
            except:
                pass

    if interactive == True:
        try:
            selection = sit.utilities.selector_list(module_list)
            used.append(selection)
            return selection, used

        except:
            raise ValueError('ERROR: Invalid selection')

    else:
        print('modules available in {}:\n'.format(DIR_MODULES))
        sit.utilities.print_list_cols(module_list)

def get_mode(expt):
    """
    Prints the available image modalities. Optionally prompts the user to
    select one.
    """
    if not DIR_DATA:
        sys.exit('ERROR: Environment variable SCRIPTUIT_DATA not defined.')

    # get a sorted list of the experiments
    expt = os.path.join(DIR_DATA, expt)
    subjects = [d for d in os.listdir(expt)
                        if os.path.isdir(os.path.join(expt, d))]

    mode_list = []
    for subj in subjects:
        subjdir = os.path.join(expt, subj)
        modes = [d for d in os.listdir(subjdir) if os.path.isdir(os.path.join(subjdir, d))]
        mode_list.extend(modes)

    mode_list = list(set(mode_list)) # keep unique entries
    try:
        selection = sit.utilities.selector_list(mode_list)
        return selection
    except:
        return None

def get_experiments(check=None):
    """
    Prints the available experiments. Optionally prompts the user to select
    one.
    """
    if sit.utilities.has_permissions(DIR_DATA) == False:
        sys.exit('ERROR: You do not have write permission to {}'.format(DIR_DATA))

    # get a sorted list of the experiments
    expt_list = [d for d in os.listdir(DIR_DATA)
                         if os.path.isdir(os.path.join(DIR_DATA, d))]

    try:
        selection = sit.utilities.selector_list(expt_list)

        # this is for the commandline usage
        if check == 'check':
            check_directories(selection)
        # this is an internal usage
        else:
            return selection

    except ValueError as ve:
        return None

def get_help(module):
    """
    Prints the module header line by line, wrapping long sections.
    """
    f = open(os.path.join(DIR_MODULES, module))
    helpfile = f.readlines()

    for line in helpfile:
        if line.startswith('#!'):
            pass
        elif line.startswith('#'):
            if len(line) < 80:
                print line.strip('#\n').lstrip(' ')
            else:
                print("\n".join(textwrap.wrap(line.strip('#').lstrip(' '))))
        else:
            break

def check_environment(mode='verbose'):
    """
    Checks manditory environent variables and exits if they defined improperly.
    """

    if DIR_DATA == None:
        sys.exit('ERROR: SCRIPTUIT_DATA variable is not set.')
    if os.path.isdir(DIR_DATA) == False:
        sys.exit('ERROR: scriptuit data directory does not exist: {}'.format(DIR_DATA))
    if sit.utilities.has_permissions(DIR_DATA) == False:
        sys.exit("ERROR: no write permissions for {}".format(DIR_DATA))

    if DIR_MODULES == None:
        sys.exit('ERROR: SCRIPTUIT_MODULES variable is not set')
    if os.path.isdir(DIR_MODULES) == False:
        sys.exit('ERROR: scriptuit modules directory does not exist: {}'.format(DIR_MODULES))

    # format uses global variables defined at top of scriptuit
    if mode == 'verbose':
        print('scriptuit configuration:\n'
              '    data directory   : {DIR_DATA}\n'
              '    module directory : {DIR_MODULES}\n'.format(
                                       DIR_DATA=DIR_DATA, DIR_MODULES=DIR_MODULES))

def check_run(expt, subj, mode, sess, run):
    """
    Returns True if RUN folder contains a NIFTI file, False otherwise.

    Usage:
        truth = check_run(experiment, subject, image modality, session, run)
    """
    # retrieve the number of files within each run folder
    dir_file = os.path.join(DIR_DATA, expt, subj, mode, sess, run)
    file_list = [d for d in os.listdir(dir_file)
                         if os.path.isfile(os.path.join(dir_file, d))]

    n_files = filter(lambda x: x[-7:] == '.nii.gz' or
                               x[-4:] == '.nii', file_list)
    n_files = len(n_files)

    return n_files

def check_mode(expt, subj, mode):
    """
    This returns each image modality's total number of runs and checks for run
    folders with missing data.
    """
    n_runs = 0

    # retrieve the sessions for a given image modality
    dir_mode = os.path.join(DIR_DATA, expt, subj, mode)
    sess_list = [d for d in os.listdir(dir_mode)
                         if os.path.isdir(os.path.join(dir_mode, d))]

    # retrieve runs for a given session
    for sess in sess_list:
        dir_sess = os.path.join(dir_mode, sess)
        run_list = [d for d in os.listdir(dir_sess)
                            if os.path.isdir(os.path.join(dir_sess, d))
                            and d[0:3] == 'RUN']

        # check for runs missing/too many NIFTI files, print warnings
        for run in run_list:
            n_files = check_run(expt, subj, mode, sess, run)
            n_runs = n_runs + n_files

            if n_files == 0:
                print('subject {}, {}, sess {}, run {} contains no NIFTI file.'.format(
                    str(subj), str(mode), str(sess), str(run)))

            if n_files > 1:
                print('subject {}, {}, sess {}, run {} contains {} NIFTI files.'.format(
                    str(subj), str(mode), str(sess), str(run), str(n_files)))
    return n_runs

def check_directories(expt):
    """
    This checks the image modalities for each subject in an experiment. It also
    reports the number of subjects with each kind of image modality.
    """
    if sit.utilities.has_permissions(os.path.join(DIR_DATA, expt)) == False:
        sys.exit('ERROR: you do not have permissions to edit this experiment.')

    # attempt to access the subjects for a given experiment
    try:
        subjects = sit.utilities.get_subj(os.path.join(DIR_DATA, expt))
    except:
        print('ERROR: experiment path is incorrect {}'.format(expt))

    # this dict will hold our count
    mode_dict = {}

    # loop through subjects, returning image modality counts, and record
    for subj in subjects:

        # retrieve the image modalities for a given subject
        dir_subj = os.path.join(DIR_DATA, expt, subj)
        mode_list = [d for d in os.listdir(dir_subj)
                             if os.path.isdir(os.path.join(dir_subj, d))]

        for mode in mode_list:
            # check inside each modality, get the number of sessions
            n_runs = check_mode(expt, subj, mode)

            # add them to the count
            try:
                mode_dict[str(mode)] = mode_dict[str(mode)] + n_runs
            except:
                mode_dict[str(mode)] = n_runs

    # print out the file counts per image modality
    print('')
    for mode_key in mode_dict:
        print('     {} NIFTIs in {}'.format(str(mode_dict[mode_key]), str(mode_key)))
    print('')

def get_loop(function, *args):
    """
    Runs the supplied function until the function no longer returns none.
    """
    var = None
    while var is None:
        var = function(*args)
    return var

def parse_master(script):
    """
    Parses submitted master script and checks formatting. Exits if any problems
    are found.
    """
    # read in master script
    masterData = [line.replace('=', ' ').rstrip() for line in open(script, 'rb')]

    # verify input script is formatted properly
    try:
        assert sit.utilities.get_line(masterData, 'DIR_MODULES'), 'ERROR: DIR_MODULES not defined.'
        assert sit.utilities.get_line(masterData, 'DIR_DATA'), 'ERROR: DIR_DATA not defined.'
        assert sit.utilities.get_line(masterData, 'DIR_EXPT'), 'ERROR: DIR_EXPT not defined.'
        assert sit.utilities.get_line(masterData, 'DATA_TYPE'), 'ERROR: DATA_TYPE not defined.'
        assert sit.utilities.get_line(masterData, 'ID'), 'ERROR: ID not defined.'
    except AssertionError as err:
        sys.exit(err)

    # verify we have at least one module defined that is also in DIR_MODULES
    moduleList = [f for f in os.listdir(DIR_MODULES)
                           if os.path.isfile(os.path.join(DIR_MODULES, f))]

    # return the modules defined in the master script that are ALSO in the
    # SCRIPTUIT_MODULES folder
    moduleList.sort()
    foundModules = []
    for line in masterData:
        for m in moduleList:
            if line.split(' ')[0] == m:
                foundModules.append(m)

    if len(foundModules) == 0:
        sys.exit('ERROR: No modules in {} are defined in {}'.format(DIR_MODULES, script))

    return masterData, foundModules

def clean(script):
    """
    Removes stages of a scriptuit pipeline.
    """
    sit.utilities.check_os()
    check_environment('quiet')
    masterData, foundModules = parse_master(script)

    DIR_MODULES = sit.utilities.get_line(masterData, 'DIR_MODULES')[0]
    DIR_DATA = sit.utilities.get_line(masterData, 'DIR_DATA')[0]
    expt = sit.utilities.get_line(masterData, 'DIR_EXPT')[0]
    mode = sit.utilities.get_line(masterData, 'DATA_TYPE')[0]
    #ID = sit.utilities.get_line(masterData, 'ID')[0] # ignore ID for now, ID to depricate?

    print('scriptuit clean: this deletes data.')

    # get the subjects
    print('clean all subjects, or some?')
    decision = get_loop(sit.utilities.selector_list, ['all', 'some'])

    if decision == 'stop':
        sys.exit()

    elif decision == 'some':
        flag = 0
        subjects = []
        subj_list = sit.utilities.get_subj(os.path.join(DIR_DATA, expt))

        while flag == 0:
            print("select subject to clean. type 'stop' to stop.")
            subj = get_loop(sit.utilities.selector_list, subj_list)

            if subj == 'stop':
                flag = 1
                continue

            subjects.append(subj)
            subj_list.remove(subj)
            print('the following subjects are slated for clean:\n')
            sit.utilities.print_list(subjects)

    elif decision == 'all':
        subjects = sit.utilities.get_subj(os.path.join(DIR_DATA, expt))
        print('all subjects slated for clean.')

    # select cutoff module
    print('select cutoff module. the outputs of this module and all downstream modules will be removed.')
    cutoff = sit.utilities.selector_list(foundModules, sort=False)
    modules = foundModules[foundModules.index(cutoff):]

    remove = []
    for module in modules:
        # parse header, get list of prefixes to remove
        header  = sit.utilities.get_header(os.path.join(DIR_MODULES, module))
        output  = sit.utilities.get_line(header, 'output:')
        others  = sit.utilities.get_line(header, 'others:')

        if output:
            remove.append(output[0])
        if others:
            remove.extend(others)

    print(remove)
    remove = (list(set(remove))) # not sure if i need this any longer...
    remove.sort()

    print('The following file types are slated to be removed:\n')
    sit.utilities.print_list(remove)

    # loop through subjects, removing specified files in the selected modality
    for prefix in remove:
        for subj in subjects:
            to_remove = [x for x in sit.utilities.find_files(
                 os.path.join(DIR_DATA, expt, subj, mode), prefix, exclude='RUN', level=2)][0]
            sit.utilities.purge(to_remove)

def render(script, output):
    """
    Renders the submitted master script.
    """
    sit.utilities.check_os()
    check_environment('quiet')
    datetime, user, f_id = sit.utilities.get_date_user()
    masterData, foundModules = parse_master(script)

    # initalize the output script
    print('rendering master script {} to output {}'.format(script, output))
    f = open(output, 'wb')
    f.write('#!/bin/bash\n\n'
            '# rendered scriptuit from {script}\n'
            '# generated: {datetime} by {user}.\n\n'
            'set -e\n\n'
            'export DIR_MODULES={DIR_MODULES}\n'
            'export DIR_DATA={DIR_DATA}\n'
            'export DIR_EXPT={DIR_EXPT}\n'
            'export DATA_TYPE={DATA_TYPE}\n'
            'export ID={ID}\n'
            'export SUB=${{1}}\n\n'
            'if [ -z ${{1}} ]; then\n'
            '    echo "Usage:"\n'
            '    echo "    $(basename ${{0}}) subject"\n'
            '    exit 1\n'
            'fi\n'.format(
                script=script,
                user=user,
                datetime=datetime,
                DIR_MODULES=sit.utilities.get_line(masterData, 'DIR_MODULES')[0],
                DIR_DATA=sit.utilities.get_line(masterData, 'DIR_DATA')[0],
                DIR_EXPT=sit.utilities.get_line(masterData, 'DIR_EXPT')[0],
                DATA_TYPE=sit.utilities.get_line(masterData, 'DATA_TYPE')[0],
                ID=sit.utilities.get_line(masterData, 'ID')[0]))

    # now, copy the contents of each scriptuit module to output, sans header
    for module in foundModules:
        for line in masterData:
            moduleBody = None # default value, if no module is found
            if line.startswith(module):

                # get module body (strip header)
                moduleBody = sit.utilities.get_body(os.path.join(DIR_MODULES, module))

                # hard code command-line variables
                moduleBody = sit.utilities.get_rendered_module(moduleBody, line)
                break

        if moduleBody:
            f.writelines(moduleBody)

    f.close()
    os.chmod(output, 0755)

def generate():
    """
    Runs the master script generator.
    """
    sit.utilities.check_os()
    check_environment('quiet')
    datetime, user, f_id = sit.utilities.get_date_user()

    # define experiment directory and check permissions
    print('\nSelect experiment:')
    expt = get_loop(get_experiments)
    if expt == 'stop':
        sys.exit()
    if sit.utilities.has_permissions(os.path.join(DIR_DATA, expt)) == False:
        sys.exit("You do not have permission to write in {}".format(os.path.join(DIR_DATA, expt)))

    print('\nexperiment: {}\nselect modality:'.format(expt))
    mode = get_loop(get_mode, expt)
    if mode == 'stop':
        sys.exit()

    # get the date and time, current user, f_id, set dir_pipe to be local
    print('\nmodality: {}\ninput ID for this run of scriptuit:'.format(mode))
    ID = raw_input('ID: ')
    ID = sit.utilities.mangle_string(ID)
    master = 'master_{}_{}.sh'.format(f_id, ID)

    print('\ninitializing master script.')
    f = open(os.path.join(DIR_DATA, expt, master), 'wb')
    f.write('#!/bin/bash\n\n'
            '# master scriptuit for {expt}: {mode}.\n'
            '# generated: {datetime} by {user}.\n\n'
            'DIR_MODULES={DIR_MODULES}\n'
            'DIR_DATA={DIR_DATA}\n'
            'DIR_EXPT={expt}\n'
            'DATA_TYPE={mode}\n'
            'ID={ID}\n\n'.format(
                DIR_MODULES=DIR_MODULES, DIR_DATA=DIR_DATA, expt=expt,
                mode=mode, ID=ID, user=user, datetime=datetime))

    p_list = []      # keeps track of our command-list
    usedModules = [] # keeps track of the modules run
    outputFiles = [] # keeps track of each module's output
    flag = False

    while not flag:

        print('\nselect the next pipeline module:')
        try:
            module, usedModules = get_modules(interactive=True, used=usedModules)
        except:
            continue

        if not module:
            continue

        if module == 'stop':
            flag = True
            continue

        # parses module, write pipeline commands
        try:
            outputFiles, command = sit.utilities.parse(os.path.join(DIR_MODULES, module), usedModules, outputFiles)
            p_list = sit.utilities.writer(f, p_list, command)
        except Exception as err:
            if len(usedModules) > 0:
                usedModules.pop() # remove last module from the used list
            print(CLEARSCREEN)
            print(err)
            continue

        print('\nCurrent pre-processing pipeline:')
        sit.utilities.print_list(p_list)
        response = raw_input("Type 'stop' to stop, else continue: ")
        if response == 'stop':
            flag = True
        print(CLEARSCREEN)

    # close the file, make executable by owner + group, and execute
    print('saving master script: {}'.format(os.path.join(DIR_DATA, expt, master)))
    f.close()
    os.chmod(os.path.join(DIR_DATA, expt, master), 0755)

if __name__ == "__main__":

    if len(sys.argv) == 2 and sys.argv[1] == 'list':
        get_modules()
    elif len(sys.argv) == 3 and sys.argv[1] == 'help':
        get_help(sys.argv[2])
    elif len(sys.argv) == 3 and sys.argv[1] == 'check' and sys.argv[2] == 'setup':
        check_environment('verbose')
    elif len(sys.argv) == 3 and sys.argv[1] == 'check' and sys.argv[2] == 'inputs':
        get_experiments(check='check')
    elif len(sys.argv) == 2 and sys.argv[1] == 'generate':
        generate()
    elif len(sys.argv) == 4 and sys.argv[1] == 'render':
        render(sys.argv[2], sys.argv[3])
    elif len(sys.argv) == 3 and sys.argv[1] == 'clean':
        clean(sys.argv[2])
    else:
        print(__doc__)

