import configargparse
import os
import signal
import torch

import onmt.opts as opts
from utils.logging import logger
import utils.distributed
from train_single import main as single_main

def main(opt):

  if opt.truncated_decoder > 0 and opt.accum_count > 1:
    raise AssertionError("BPTT is not compatible with -accum > 1")

  if opt.gpuid:
    raise AssertionError("gpuid is deprecated \
          see world_size and gpu_ranks")

  nb_gpu = len(opt.gpu_ranks)

  if opt.world_size > 1:
    mp = torch.multiprocessing.get_context('spawn')
    # Create a thread to listen for errors in the child processes.
    error_queue = mp.SimpleQueue()
    error_handler = ErrorHandler(error_queue)
    # Train with multiprocessing.
    procs = []
    for device_id in range(nb_gpu):
      procs.append(mp.Process(target=run, args=(
          opt, device_id, error_queue, ), daemon=True))
      procs[device_id].start()
      logger.info(" Starting process pid: %d  " % procs[device_id].pid)
      error_handler.add_child(procs[device_id].pid)
    for p in procs:
      p.join()

  elif nb_gpu == 1:  # case 1 GPU only
    single_main(opt, 0)
  else:   # case only CPU
    single_main(opt, -1)



def run(opt, device_id, error_queue):
  """ run process """
  try:
    gpu_rank = utils.distributed.multi_init(opt, device_id)
    if gpu_rank != opt.gpu_ranks[device_id]:
      raise AssertionError("An error occurred in \
        Distributed initialization")
    single_main(opt, device_id)
  except KeyboardInterrupt:
    pass  # killed by parent, do nothing
  except Exception:
    # propagate exception to parent process, keeping original traceback
    import traceback
    error_queue.put((opt.gpu_ranks[device_id], traceback.format_exc()))

class ErrorHandler(object):
  """A class that listens for exceptions in children processes and propagates
  the tracebacks to the parent process."""

  def __init__(self, error_queue):
    """ init error handler """
    import threading
    self.error_queue = error_queue
    self.children_pids = []
    self.error_thread = threading.Thread(
        target=self.error_listener, daemon=True)
    self.error_thread.start()
    signal.signal(signal.SIGUSR1, self.signal_handler)

  def add_child(self, pid):
    """ error handler """
    self.children_pids.append(pid)

  def error_listener(self):
    """ error listener """
    (rank, original_trace) = self.error_queue.get()
    self.error_queue.put((rank, original_trace))
    os.kill(os.getpid(), signal.SIGUSR1)

  def signal_handler(self):
    """ signal handler """
    for pid in self.children_pids:
        os.kill(pid, signal.SIGINT)  # kill children processes
    (_, original_trace) = self.error_queue.get()
    msg = """\n\n-- Tracebacks above this line can probably
             be ignored --\n\n"""
    msg += original_trace
    raise Exception(msg)

if __name__ == "__main__":
  parser = configargparse.ArgumentParser(
    description='train.py',
    config_file_parser_class=configargparse.YAMLConfigFileParser,
    formatter_class=configargparse.ArgumentDefaultsHelpFormatter)

  opts.config_opts(parser)
  #opts.add_md_help_argument(parser)
  opts.model_opts(parser)
  opts.train_opts(parser)

  opt = parser.parse_args()
  main(opt)