#! /usr/bin/python3

import bisect as BS
import curses as CU
import optparse as OP
import os as OS
import random as RND
import resource as RSC
import subprocess as SUB
import sys as SYS

op = OP.OptionParser\
       (usage = "%prog [-t] [-d FILE] [-j JOBS] [-n SAMPLES] [-r REF] VAR ...")
op.disable_interspersed_args()
for short, long, kw in \
    [("-d", "--dictionary",
      dict(type = "string", metavar = "DICT",
           dest = "dict", default = "/usr/share/dict/words",
           help = "word list to process")),
     ("-j", "--jobs",
      dict(type = "int", metavar = "JOBS",
           dest = "njobs", default = 1,
           help = "number of jobs to run concurrently")),
     ("-n", "--samples",
      dict(type = "int", metavar = "SAMPLES",
           dest = "nsamples", default = 5,
           help = "number of samples for each variant")),
     ("-r", "--reference",
      dict(type = "string", metavar = "REF",
           dest = "ref", default = None,
           help = "reference output data expected")),
     ("-t", "--terminal",
      dict(action = "store_true", dest = "terminal", default = None,
           help = "show progress display"))]:
  op.add_option(short, long, **kw)
OPTS, ARGS = op.parse_args()

REF = OPTS.ref
if REF is None:
  REF = "REF.%s" % OS.path.basename(OPTS.dict)
  if not OS.path.exists(REF): REF = None

IDLE = 0
LIVE = 1
DONE = 2

class BasicVariant (object):

  def __init__(me, name, row):
    me.name = name
    me.samples = []
    me.row = row
    me._prog = "./chain.%s" % name
    me.state = IDLE

  @property
  def n(me):
    return len(me.samples)

  def start(me, out):
    SYS.stdin.close()
    fd = OS.open(OPTS.dict, OS.O_RDONLY)
    if fd != 0: OS.dup2(fd, 0); OS.close(fd)
    SYS.stdout.close()
    fd = OS.open(out, OS.O_WRONLY | OS.O_CREAT | OS.O_TRUNC, 0o666)
    if fd != 1: OS.dup2(fd, 1); OS.close(fd)
    OS.execv(me._prog, [me._prog])

  def end(me, time):
    me.samples.append(time)
    return time

class LispVariant (BasicVariant):
  def __init__(me, name, *args, **kw):
    super(LispVariant, me).__init__(name, *args, **kw)
    me._time = "time.%s" % name
    me._lisp = name[5:]
    me._prog = "chain.lisp"

  def start(me, out):
    SYS.stdout.close()
    fd = OS.open(out, OS.O_WRONLY | OS.O_CREAT | OS.O_TRUNC, 0o666)
    if fd != 1: OS.dup2(fd, 1); OS.close(fd)
    OS.execvp("runlisp", ["runlisp", "-L%s" % me._lisp,
                          me._prog, "-T%s" % me._time, OPTS.dict])

  def end(me, time):
    #with open(me._time) as f: time = 1000*float(f.read())
    OS.unlink(me._time)
    me.samples.append(time)
    return time

def strwait(st):
  if OS.WIFEXITED(st): return "rc = %d" % OS.WEXITSTATUS(st)
  elif OS.WIFSIGNALED(st): return "signal = %d" % OS.WTERMSIG(st)
  else: return "?st = %d" % st

class State (object):

  def __init__(me):
    me.vars = []
    me.idle = []
    me.done = []
    me.live = {}
    me.bad = False

    me._nextrow = 0
    me._dpy = None

  def _setdpy(me, dpy):
    me._dpy = dpy

  def addvar(me, var):
    if arg.startswith("lisp-"): var = LispVariant(arg, me._nextrow)
    else: var = BasicVariant(arg, me._nextrow)
    me.vars.append(var); me.idle.append(var); me._nextrow += 1

  def cycle(me):
    if me.idle and not me.bad and len(me.live) < OPTS.njobs:
      n = len(me.idle)
      k = RND.randrange(n); pick = me.idle[k]; assert pick.state == IDLE
      me.idle[k] = me.idle[n - 1]; me.idle.pop()
      kid = OS.fork()
      if not kid: pick.start("chain.%s.out" % pick.name)
      #me._dpy.msg(";; start %s\n" % pick.name)
      me.live[kid] = pick; pick.state = LIVE
      me._dpy.paint_var(pick)
      return True
    elif me.live:
      kid, st, ru = OS.wait3(0)
      try: var = me.live.pop(kid)
      except KeyError:
        me._dpy.msg("ignoring unexpected child %d (%s)" %
                    (kid, strwait(st)))
        return True
      assert var.state == LIVE
      if not (OS.WIFEXITED(st) and OS.WEXITSTATUS(st) == 0):
        me._dpy.msg("chain.%s failed (%s)" % (var.name, strwait(st)))
        me.bad = True
        return True
      out = "chain.%s.out" % var.name
      if REF is not None:
        if SUB.call(["./chkref", out, REF]):
          me._dpy.msg("chain.%s produced incorrect output" % var.name)
          me.bad = True
          return True
      OS.unlink(out)
      time = var.end(1000*(ru.ru_utime + ru.ru_stime))
      #me._dpy.msg(";; end %s, time = %.3f s" % (var.name, time))
      if var.n < OPTS.nsamples:
        me.idle.append(var); var.state = IDLE
      else:
        me.done.append(var); var.state = DONE
        #me._dpy.msg(";; done %s\n" % var.name)
      me._dpy.paint_var(var)
      return True
    else:
      return False

if OPTS.terminal or SYS.stdout.isatty():

  class Display (object):

    def __init__(me, state):
      me._st = state
      me._st._setdpy(me)
      namewd = 16
      for var in me._st.vars:
        if len(var.name) >= namewd: namewd = len(var.name) + 1
      me._a_idle = CU.A_NORMAL
      me._a_live = CU.A_BOLD
      me._namewd = namewd
      me._msgs = []

    def __enter__(me):

      if not SYS.stdout.isatty():
        ttyfd = OS.open("/dev/tty", OS.O_RDWR)
        SYS.stdout = OS.fdopen(OS.dup(1), "w")
        OS.dup2(ttyfd, 1)
        OS.close(ttyfd)

      me._scr = CU.initscr()
      me._resize()
      CU.start_color()
      CU.use_default_colors()
      if CU.COLORS and CU.COLOR_PAIRS >= 3:
        CU.init_pair(1, CU.COLOR_GREEN, -1)
        me._a_done = CU.color_pair(1)
        CU.init_pair(2, CU.COLOR_BLACK, CU.COLOR_GREEN)
        CU.init_pair(3, CU.COLOR_BLACK, CU.COLOR_YELLOW)
        me._a_bar_left = CU.color_pair(2)
        me._a_bar_right = CU.color_pair(3)
      else:
        me._a_live = CU.A_DIM
        me._a_bar_left = CU.A_REVERSE
        me._a_bar_right = None
      me._msgrow = min(len(me._st.vars) + 1, me._ht - 1)
      me._nmsgs = me._ht - me._msgrow
      me._scr.setscrreg(me._msgrow, me._ht - 1)
      me._scr.scrollok(True)
      me._repaint()

    def __exit__(me, exval, exty, tb):
      CU.endwin()
      for m in me._msgs: SYS.stderr.write(m + "\n")

    def _resize(me):
      me._ht, me._wd = me._scr.getmaxyx()
      barwd = me._wd - me._namewd
      if OPTS.nsamples <= barwd: me._barsc = 1
      else: me._barsc = float(barwd)/OPTS.nsamples

    def paint_var(me, var, refresh = True):
      if var.row >= me._ht - 2: return
      rbar = int(me._barsc*OPTS.nsamples + 0.5)
      if var.state == LIVE: attr = me._a_live
      elif var.state == DONE: attr = me._a_done
      else: attr = me._a_idle
      me._scr.addstr(var.row, 0, var.name, attr)
      mbar = int(me._barsc*len(var.samples))
      me._scr.attrset(me._a_bar_left)
      me._scr.hline(var.row, me._namewd, " ", mbar)
      if me._a_bar_right is not None:
        me._scr.attrset(me._a_bar_right)
        me._scr.hline(var.row, me._namewd + mbar, " ", rbar - mbar)
        me._scr.move(var.row, me._namewd + rbar)
      else:
        me._scr.move(var.row, me._namewd + mbar)
      me._scr.attrset(0)
      if var.samples:
        avg = "%.3f ms" % (sum(var.samples)/len(var.samples))
        if me._wd - (me._namewd + rbar) > len(avg):
          me._scr.addstr(var.row, me._namewd + rbar + 1, avg)
      me._scr.clrtoeol()
      if refresh: me._scr.refresh()

    def msg(me, msg):
      n = len(me._msgs)
      if n >= me._nmsgs:
        me._scr.scroll(1)
      me._scr.addstr(min(me._msgrow + n, me._ht - 1), 0, msg)
      me._msgs.append(msg)
      me._scr.refresh()

    def _repaint(me):
      me._scr.erase()
      for var in me._st.vars: me.paint_var(var, refresh = False)
      n = len(me._msgs); y = me._msgrow
      for i in range(max(0, n - me._nmsgs), n):
        me._scr.addstr(y, 0, me._msg[i]); y += 1
      me._scr.refresh()

else:

  class Display (object):

    def __init__(me, state):
      me._st = state
      me._st._setdpy(me)

    def __enter__(me):
      pass
    def __exit__(me, exval, exty, tb):
      pass

    def paint_var(me, var):
      if var.state == IDLE:
        SYS.stderr.write(";; end %s, time = %.3f ms\n" %
                         (var.name, var.samples[-1]))
      elif var.state == LIVE:
        SYS.stderr.write(";; start %s\n" % var.name)
      elif var.state == DONE:
        SYS.stderr.write(";; %s all done\n" % var.name)

    def msg(me, msg):
      SYS.stderr.write("%s\n" % msg)

RND.seed()
STATE = State()
for arg in ARGS: STATE.addvar(arg)
with Display(STATE):
  while STATE.cycle(): pass
if STATE.bad: SYS.exit(2)

def quartile(i, x):
  n = len(x)
  k = i*n
  j = k//4
  if not k%4: return x[j]
  else: return (x[j] + x[j + 1])/2

def stats(x):
  q1, mid, q3 = quartile(1, x), quartile(2, x), quartile(3, x)
  iqr = q3 - q1
  locut = BS.bisect_left(x, q1 - 1.5*iqr); lo = x[locut]
  hicut = BS.bisect_right(x, q3 + 1.5*iqr); hi = x[hicut - 1]
  avg = sum(x)/len(x)
  return lo, q1, mid, avg, q3, hi

done = STATE.vars
done.sort(key = lambda var: var.name)
n = OPTS.nsamples - 1
for var in done:
  x = [xi for xi in var.samples]; x.sort()
  lo, q1, mid, avg, q3, hi = stats(x)
  print("%-16s %s; %s" %
        (var.name,
         " ".join("%.3f" % xi for xi in var.samples),
         " ".join("%.3f" % xi for xi in [lo, q1, mid, avg, q3, hi])))