#=============================================================================
#
#  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'

import imp
import os
import sys
import time
import traceback
import shutil

import girder_client
print 'girder_client', girder_client.__version__
import requests
import smtk
import smtk.attribute

import cumulusclient
reload(cumulusclient)  # for dev
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:
    # Verify simulation files
    check_file(scope.model_path, 'Cannot find model file at %s')

    extension_lookup = {
      'tem3p-eigen': 'tem3p',
      'tem3p-structure': 'tem3p',
      'tem3p-harmonic': 'tem3p',
      'tem3p-thermal-linear': 'tem3p',
      'tem3p-thermal-nonlinear': 'tem3p'
    }
    for solver in scope.solver_list:
      if solver == 'acdtool':
        continue
      file_ext = extension_lookup.get(solver, solver)
      ace3p_filename = '%s.%s' % (scope.output_file_prefix, file_ext)
      path = os.path.join(scope.output_folder, ace3p_filename)
      check_file(path, 'Cannot find simulation input file at %s')

    # Start NERSC session
    login_nersc(scope, sim_item)

    # 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 'ERROR', err.responseText
    raise
  except Exception as ex:
    print 'Exception', ex
    traceback.print_exc()
    raise
  finally:
    #release_resources(scope)
    if scope.nersc:
      scope.nersc.logout()

  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):
  '''Logs into NERSC and gets session id

  '''
  scope.newt_sessionid = None

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

  nersc_url = 'https://newt.nersc.gov/newt'
  scope.nersc = NewtClient(nersc_url)
  r = scope.nersc.login(username, password)
  scope.newt_sessionid = scope.nersc.get_sessionid()
  print 'newt_sessionid', scope.newt_sessionid

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

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

  # Process acdtool separately
  root, ext = os.path.splitext(scope.model_file)
  if scope.solver_list[0] == 'acdtool':
    args = get_acdtool_args(scope)
    convert_command = 'srun -n 1 %s/acdtool %s' % (binary_folder, args)
    command_list.append(convert_command)
  elif '.gen' == ext:
    # If model file is .gen, convert to ncdf
    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)

  code_lookup = {
    'tem3p-eigen': 'tem3p',
    'tem3p-structure': 'tem3p',
    'tem3p-harmonic': 'tem3p',
    'tem3p-thermal-linear': 'tem3p',
    'tem3p-thermal-nonlinear': 'tem3p'
  }

  # Add the solver commands(s)
  for solver in scope.solver_list:
    if 'acdtool' == solver:
      continue
    code = code_lookup.get(solver, solver)
    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()

  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

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

  return concrete_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

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

  return concrete_item.value(0)
