#!/usr/bin/env python

import os
import sys
import argparse
import subprocess
import spikeforest as sf
import batcho
import time

# This can be important for some of the jobs in certain situations
os.environ['DISPLAY']=''

def main():
  parser = argparse.ArgumentParser(description = 'Run a spikeforest processing batch')
  parser.add_argument('--run_prefix',help='Prefix to run command', default='')
  parser.add_argument('--clear',help='First clear the jobs so that all jobs are run.',action="store_true")
  parser.add_argument('--mlpr_force_run',help='Force run the MountainLab processors',action="store_true")
  parser.add_argument('--job_index',help='Only run one job and do not assemble the results',default='')
  parser.add_argument('--parallel',help='Number of times to simultaneously do the run command.', default='1')
  parser.add_argument('--repeat',help='Continuously repeat the processing with a delay between successive runs',action="store_true")
  parser.add_argument('--delay',help='Number of seconds to delay between runs (to be used with --repeat)',default='5')
  parser.add_argument('batch_names',help='The names of the batches to run',nargs='*')

  args = parser.parse_args()

  run_prefix=args.run_prefix
  if run_prefix:
    run_prefix=run_prefix+' '

  prepare_opts=[]
  run_opts=[]
  assemble_opts=[]
  if args.clear:
    prepare_opts.append('--clear')
  if args.mlpr_force_run:
    run_opts.append('--mlpr_force_run')
  if args.job_index:
    prepare_opts.append('--job_index {}'.format(args.job_index))
    run_opts.append('--job_index {}'.format(args.job_index))
  prepare_opts=' '.join(prepare_opts)
  run_opts=' '.join(run_opts)
  assemble_opts=' '.join(assemble_opts)
  if args.job_index:
    job_index=int(args.job_index)
  else:
    job_index=None

  dirname=os.path.dirname(os.path.realpath(__file__))

  if len(args.batch_names)==0:
    raise Exception('At least one batch name is needed.')

  print('Starting sf_run_batch...')
  while True:

    for batch_name in args.batch_names:

      print('BATCH: '+batch_name)
      if (args.clear) or (_check_run_needed(batch_name, job_index)):

        ## PREPARE
        _system_call(dirname+'/sf_run_batch_command prepare {} {}'.format(prepare_opts, batch_name))

        ## RUN
        def do_run(num=0):
          _system_call(run_prefix+dirname+'/sf_run_batch_command run {} {}'.format(run_opts, batch_name))

        num_parallel=int(args.parallel)
        if num_parallel<=1:
          do_run()
        else:
          from multiprocessing import Pool
          pool=Pool()
          pool.map(do_run, list(range(num_parallel)))

        ## ASSEMBLE
        if job_index is not None:
          print('Not assembling because --job_index option was used.')
        else:
          _system_call(dirname+'/sf_run_batch_command assemble {} {}'.format(assemble_opts, batch_name))

    if not args.repeat:
      break

    delay_sec=int(args.delay)
    print('Waiting {} seconds before repeating'.format(delay_sec))
    time.sleep(delay_sec)

def _check_run_needed(batch_name, job_index):
  sf.kbucketConfigRemote(name='spikeforest1-readonly')
  print('Checking if batch run is needed for {}...'.format(batch_name))
  try:
    statuses=batcho.get_batch_job_statuses(batch_name=batch_name, job_index=job_index)
  except:
    statuses=None
  if not statuses:
    print('Unable to find statuses for batch: '+batch_name)
    return False
  print('Found {} jobs...'.format(len(statuses)))
  for status in statuses:
    if status['status'] != 'finished':
      return True
  print('Batch already completed.')
  return False

def _system_call(cmd):
  print('EXECUTING: '+cmd)
  p = subprocess.Popen(cmd, shell=True, stderr=subprocess.PIPE)
   
  while True:
      out = p.stderr.read(1)
      if not out:
        if p.poll() != None:
            break
      if out:
          sys.stdout.write(out.decode())
          sys.stdout.flush()

  if p.returncode != 0:
    raise Exception('Non-zero exit code for command: '+cmd)

if __name__== "__main__":
  main()
  
