"""Functions to support job submission using NEWT interface."""

import os

import smtk
import smtk.attribute

RECURSIVE = smtk.attribute.SearchStyle.RECURSIVE


def create_slurm_commands(scope, sim_item):
    """Returns SLURM script as a list of strings.

    sim_item is the NERSCSimulation GroupItem.
    """
    # Get machine info
    number_of_nodes = get_integer(sim_item, 'NumberOfNodes')
    number_of_tasks = get_integer(sim_item, 'NumberOfTasks')
    total_number_of_tasks = number_of_nodes * number_of_tasks

    # CPU per task depends on the machine
    machine = get_string(sim_item, 'Machine')
    if machine == 'edison':
        cpu_per_task = 48 // number_of_tasks    # (integer division)
    elif machine == 'cori':
        cpu_per_task = 64 // number_of_tasks
    else:
        raise Exception('Unrecognized machine name: %s' % machine)

    command_list = list()    # return value
    command_list.append('#!/bin/bash')

    job_name = get_string(sim_item, 'JobName')
    command_list.append('#SBATCH --job-name={}'.format(job_name))

    qos = get_string(sim_item, 'Queue')
    command_list.append('#SBATCH --qos={}'.format(qos))

    timeout = get_integer(sim_item, 'Timeout')
    command_list.append('#SBATCH --time={}'.format(timeout))

    command_list.append('#SBATCH --nodes={}'.format(number_of_nodes))
    command_list.append('#SBATCH --constraint=haswell')

    # Put placeholder in for workding directory
    command_list.append('#SBATCH --chdir=$JOB_DIRECTORY')

    # Should this be added as a user option?
    command_list.append('#SBATCH --partition=debug')

    accountItem = sim_item.find('NERSCRepository')
    command_list.append('#SBATCH --account={}'.format(accountItem.value()))

    # Set modules
    command_list.append('module -s swap PrgEnv-intel PrgEnv-gnu')
    command_list.append('module load gsl')

    command_list.append('ulimit -s unlimited')    # stack size
    binary_folder = '/project/projectdirs/ace3p/{}'.format(machine)

    # Check for scope.symlink first
    if scope.symlink:
        command = 'ln -f -s %s .' % scope.symlink
        command_list.append(command)

    # print('CHECKPOINT', scope.solver_list, scope.categories, scope.output_file_prefix)
    # Process acdtool separately
    if scope.solver_list[0] == 'Rf-Postprocess':
        # Get name of results directory from sim atts
        rfpost_att = scope.sim_atts.findAttribute('RfPostprocess')
        result_dir_item = rfpost_att.findString('ResultDir')
        result_dir = result_dir_item.value()
        print('result_dir', result_dir)

        # Get full path to (previously-run) simulation job (folder)
        data_att = scope.sim_atts.findAttribute('InputData')
        data_folder = data_att.findString('NERSCDataDirectory').value().rstrip('/')
        scope.nersc_input_folder = data_folder
        print('Set scope.nersc_input_folder to', scope.nersc_input_folder)

        # Add symbolic link
        ln_command = 'ln -s {}/{} .'.format(data_folder, result_dir)
        command_list.append(ln_command)

        input_filename = '{}.rfpost'.format(scope.output_file_prefix)
        srun_command = 'srun -n 1 {}/acdtool postprocess rf {}'.format(binary_folder, input_filename)
        command_list.append(srun_command)
        return command_list
    else:
        # Check for optional data directory
        data_att = scope.sim_atts.findAttribute('InputData')
        data_item = data_att.itemAtPath('Source/NERSCDirectory')
        if data_item and data_item.isEnabled():
            data_folder = data_item.value().rstrip('/')
            scope.nersc_input_folder = data_folder
            print('Set scope.nersc_input_folder to', scope.nersc_input_folder)

    # Check if we need to convert model file
    root, ext = os.path.splitext(scope.model_file)
    if '.gen' == ext:
        print('Adding command to convert %s to netcdf' % scope.model_file)
        convert_command = \
            ('srun -n 1 %s/acdtool meshconvert %s') % (binary_folder, scope.model_file)
        command_list.append(convert_command)

    # Add the solver commands(s)
    for solver in scope.solver_list:
        if 'acdtool' == solver:
            continue
        code = solver.lower()
        ace3p_filename = '%s.%s' % (scope.output_file_prefix, code)
        results_folder = get_string(sim_item, 'ResultsDirectory')
        sim_command = 'srun -n %s -c %s %s/%s %s %s' % \
            (total_number_of_tasks, cpu_per_task, binary_folder, code, ace3p_filename, results_folder)
        command_list.append(sim_command)

    return command_list

# ---------------------------------------------------------------------
def get_integer(group_item, name, style=RECURSIVE):
    '''Looks for IntItem contained by group.

    Returns either integer value or None if not found
    '''
    item = group_item.find(name, style)
    if not item:
        print('WARNING: item \"%s\" not found' % name)
        return None

    if not item.isEnabled():
        return None

    if item.type() != smtk.attribute.Item.IntType:
        print('WARNING: item \"%s\" not an integer item' % name)
        return None

    return item.value(0)

# ---------------------------------------------------------------------
def get_string(group_item, name, style=RECURSIVE):
    '''Looks for StringItem contained by group.

    Returns either string or None if not found
    '''
    item = group_item.find(name, style)
    if not item:
        print('WARNING: item \"%s\" not found' % name)
        return None

    if not item.isEnabled():
        return None

    if item.type() != smtk.attribute.Item.StringType:
        print('WARNING: item \"%s\" not a string item' % name)
        return None

    return item.value(0)
