#=============================================================================
#
#  Copyright (c) Kitware, Inc.
#  All rights reserved.
#  See LICENSE.txt for details.
#
#  This software is distributed WITHOUT ANY WARRANTY; without even
#  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
#  PURPOSE.  See the above copyright notice for more information.
#
#=============================================================================
"""
Export functions for submitting jobs to NERSC
"""
print('Loading nersc')

from imp import reload
import os
import sys
import time
import traceback
import shutil

try:
  import girder_client
  print('girder_client', girder_client.__version__)
except ImportError:
  print('ImportError girder_client')
  print('sys.path:')
  for p in sys.path:
    print('  ', p)
  raise

import smtk
import smtk.attribute

from . import cumulusclient
reload(cumulusclient)  # for dev
from . import newtclient
reload(newtclient)  # for dev
from .cumulusclient import CumulusClient
from .newtclient import NewtClient

# ---------------------------------------------------------------------
def submit_ace3p(scope, sim_item):
  '''Submits ACE3_ job to NERSC via cumulus server

  This is the module entry point.
  Returns boolean indicating success.
  '''
  ok = True   # return value
  scope.cumulus = None
  scope.nersc = None

  try:
    # Check solver_list
    if len(scope.solver_list) != 1:
      n = len(scope.solver_list)
      msg = 'Internal Error: solver list length wrong; should be 1 not {}'.format(n)
      raise RuntimeError(msg)
    solver = scope.solver_list[0].lower()
    print('solver', solver)

    # Verify upload files
    for path in scope.files_to_upload:
      check_file(path, 'Cannot find simulation input file at %s')

    # Start NERSC session
    scope.newt_sessionid = None
    sessionid = login_nersc(scope, sim_item)
    print('newt_sessionid', sessionid)
    if sessionid is None or sessionid == "":
      raise RuntimeError('ERROR: Unable to get valid session id')
    scope.newt_sessionid = sessionid

    # Initialize CumulusClient
    scope.cumulus = create_cumulus_client(scope, sim_item)

    # Create cluster
    machine = get_string(sim_item, 'Machine')
    print('machine', machine)
    if not machine:
      raise Exception('Machine name not specified')
    scope.cumulus.create_cluster(machine)

    # Create run script
    commands = create_slurm_commands(scope, sim_item)
    print('commands', commands)
    script_name = '-'.join(scope.solver_list)
    scope.cumulus.create_slurm_script(script_name, commands)

    # Create job and upload files
    create_job(scope, solver, sim_item)
    scope.cumulus.upload_inputs(scope.files_to_upload, scope.folders_to_upload)

    # Submit job
    submit_job(scope, sim_item)
    print('Submitted %s job, id %s' % (solver, scope.cumulus.job_id()))
  except girder_client.HttpError as err:
    print('HTTP ERROR', err.responseText)
    raise
  except Exception as ex:
    print('Exception', ex)
    traceback.print_exc()
    raise
  return ok

# ---------------------------------------------------------------------
def release_resources(scope):
  '''
  '''
  print('Releasing resources')
  if scope.cumulus:
    scope.cumulus.release_resources()

  if scope.nersc:
    scope.nersc.logout()

# ---------------------------------------------------------------------
def login_nersc(scope, sim_item):
  '''Returns NEWT session id, logging into NERSC if specified

  Note that this method must *also* initialize NewtClient instance
  '''
  nersc_url = 'https://newt.nersc.gov/newt'

  # Check if sesson id provided
  credentials_item = sim_item.find('NERSCCredentials')
  credentials_type = credentials_item.value()
  if credentials_type == 'newt_sessionid':
    session_item = credentials_item.find('NEWTSessionId')
    session_id = session_item.value()
    scope.nersc = NewtClient(nersc_url, session_id)
    return session_id

  elif credentials_type != 'login':
    raise RuntimeError('ERROR - Unrecognized credentials_type', credentials_type)

  # Check user inputs
  username = get_string(sim_item, 'NERSCAccountName')
  print('username', username)
  # if not username:
  #   raise Exception('ERROR: NERSC account name not specified')

  password = get_string(sim_item, 'NERSCAccountPassword')
  print('password length', len(password))
  # if not password:
  #   raise Exception('ERROR: NERSC account password not specified')

  mfa = get_string(sim_item, 'NERSCMultfactorToken')

  scope.nersc = NewtClient(nersc_url)
  r = scope.nersc.login(username, password+mfa)
  return scope.nersc.get_sessionid()

# ---------------------------------------------------------------------
def create_cumulus_client(scope, sim_item):
  '''Instantiates Cumulus client

  '''
  cumulus_item = sim_item.find('CumulusHost')
  # cumulus_item = smtk.attribute.to_concrete(item)
  cumulus_host = cumulus_item.value(0)
  if not cumulus_host:
    raise Exception('ERROR: Cumulus host not specified')

  return CumulusClient(cumulus_host, scope.newt_sessionid)

# ---------------------------------------------------------------------
def create_slurm_commands(scope, sim_item):
  '''Creates and returns list of commands to run on the remote machine

  '''
  command_list = list()  # return 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/{{machine}}'

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

  # 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)

  # 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

  # 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)
    sim_command = 'srun -n %s -c %s %s/%s %s' % \
      (total_number_of_tasks, cpu_per_task, binary_folder, code, ace3p_filename)
    command_list.append(sim_command)

  return command_list

# ---------------------------------------------------------------------
def create_job(scope, solver, sim_item):
  '''
  '''
  job_name = get_string(sim_item, 'JobName')
  if not job_name:
    raise Exception('No Job Name specified -- cannot submit to NERSC')
  # Next line crashes modelbuilder, because there is no TailFile item any more
  # tail = get_string(sim_item, 'TailFile')
  # scope.cumulus.create_job(job_name, tail=tail)
  scope.cumulus.create_job(job_name)

  # Add cmb-specific metadata
  cmb_data = dict()
  cmb_data['solver'] = solver
  cmb_data['notes'] = ''
  if sim_item.definition().version() > 0:
    notes = get_string(sim_item, 'JobNotes')
    cmb_data['notes'] = notes

  # Number of nodes
  number_of_nodes = get_integer(sim_item, 'NumberOfNodes')
  cmb_data['numberOfNodes'] = number_of_nodes

  # Total number of cores (1 core per task times number of nodes)
  number_of_tasks = get_integer(sim_item, 'NumberOfTasks')
  cmb_data['numberOfCores'] = number_of_nodes * number_of_tasks

  # Time stamp (seconds since epoci)
  cmb_data['startTimeStamp'] = time.time()
  scope.cumulus.set_job_metadata(cmb_data)


# ---------------------------------------------------------------------
def submit_job(scope, sim_item):
  '''Run the job

  Call this after create_job()
  '''
  # Get inputs
  job_name = get_string(sim_item, 'JobName')
  machine = get_string(sim_item, 'Machine')
  number_of_nodes = get_integer(sim_item, 'NumberOfNodes')
  project_repo = get_string(sim_item, 'NERSCRepository')

  # Parse Queue input into separate queue & qos
  queue = 'debug'  # default
  qos = None
  queue_string = get_string(sim_item, 'Queue')
  if queue_string != 'debug':
    queue = 'regular'
    qos = queue_string

  # Job directory has several options
  job_output_dir = ''
  jobdir_item = sim_item.find('JobDirectory')
  if jobdir_item.type() == smtk.attribute.Item.StringType:
    job_output_dir = jobdir_item.value()  # backward compatibility
  elif jobdir_item.type() == smtk.attribute.Item.GroupType:
    filesystem_item = jobdir_item.find('FileSystem')
    if filesystem_item.value() == 'scratch':
      scratch_dir = scope.nersc.get_scratch_folder(machine)
      print('Scratch filesystem root directory:', scratch_dir)
      folder_item = filesystem_item.findChild(
        'SubFolder', smtk.attribute.ACTIVE_CHILDREN)
      folder = folder_item.value()
      job_output_dir = '%s/%s' % (scratch_dir, folder)

      # Optionally append job name to the path
      add_jobname_item = filesystem_item.findChild(
        'AppendJobNameFolder', smtk.attribute.ACTIVE_CHILDREN)
      if add_jobname_item and add_jobname_item.isEnabled():
        job_output_dir = '%s/%s' % (job_output_dir, job_name)
    elif filesystem_item.value() == 'other':
      path_item = filesystem_item.findChild(
        'FullPath', smtk.attribute.ACTIVE_CHILDREN)
      job_output_dir = path_item.value()

  # Pass back remote path including "job id" subfolder
  scope.nersc_job_folder = job_output_dir + '/' + scope.cumulus.job_id()

  timeout_minutes = get_integer(sim_item, 'Timeout')
  scope.cumulus.submit_job(
    machine, project_repo, job_output_dir, timeout_minutes, \
    queue=queue, qos=qos, number_of_nodes=number_of_nodes)

# ---------------------------------------------------------------------
def get_acdtool_args(scope):
  '''Generates command and arguments to pass to acdtool process

  Returns a string with the formatted arguments
  '''
  # Need the Analysis attribute item
  solver_item = scope.export_att.findString('Analysis')
  task_item = solver_item.findChild('AcdtoolTask', smtk.attribute.ACTIVE_CHILDREN)
  task = task_item.value(0)

  # Get extension for current model file
  root, ext = os.path.splitext(scope.model_file)

  if task == 'meshconvert':
    if ext == '.ncdf':
      raise Exception('Input model is already in ncdf format')
    return 'meshconvert %s' % scope.model_file

  elif task == 'mesh':
    if ext != '.ncdf':
      raise Exception('Run meshconvert to generate .ncdf')
    subtask_item = task_item.findChild('AcdtoolMesh', smtk.attribute.ACTIVE_CHILDREN)
    subtask = subtask_item.value(0)
    arg_list = ['mesh', subtask, scope.model_file]
    if subtask == 'fix':
      # Append output filename
      output_item = subtask_item.findChild('OutputFilename', smtk.attribute.ACTIVE_CHILDREN)
      output_filename = output_item.value(0)
      arg_list.append(output_filename)
    arg_string = ' '.join(arg_list)
    return arg_string

  elif task == 'postprocess':
    raise Exception('Support for Acttool postprocess not implemented')


  # (else)
  raise Exception('Unrecognized Acdtool Task %s' % task)


# ---------------------------------------------------------------------
def check_file(path, error_message_format=None):
  '''Confirms that file exists at given path

  Throws an exception if file not found
  '''
  if not error_message_format:
    error_message_format = 'Cannot find file at %s'
  if not os.path.isfile(path):
    raise Exception(error_message_format % path)

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

  Returns either integer value or None if not found
  '''
  item = group_item.find(name)
  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):
  '''Looks for StringItem contained by group.

  Returns either string or None if not found
  '''
  item = group_item.find(name)
  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)
