#!/usr/bin/env python

import sys
import os

import time

import shelve
from   progressbar import *
from   collections import defaultdict
from   optparse    import OptionParser

import numpy as np

import interp.bootstrap

import logging
log = logging.getLogger("interp")

from interp.cluster import QueueManager, get_qs

if __name__ == '__main__':
  parser = OptionParser(usage = "usage: %s [options] <server> <expected participants> <interp count>")

  parser.add_option("-l", "--last-time",
    action="store_true", dest="last", default=False,
    help="when finished, send shutdown signal to <expected participants> nodes (default: %default)")

  parser.add_option('-p', '--port',
    type="int", dest="port", default=6666,
    help="specify the port to use on the server (default: %default)")

  parser.add_option("-o", "--order",
    type="int", dest="order", default=2,
    help="order of interpolation (default: %default)")

  parser.add_option("-e", "--extra-points",
    type="int", dest="extra", default=3,
    help="number of extra points (default: %default)")

  parser.add_option('-s', '--shelve',
    type="str", dest="shelvename", default=os.path.expanduser('~/interp.shelve'),
    help="shelve output file (default: %default)")

  (options, args) = parser.parse_args()
  if len(args) != 3:
    parser.print_usage()
    sys.exit(1)

  server = args[0]
  expected_participants, count = (int(i) for i in args[1:])

  m = QueueManager(address=(server, options.port), authkey='asdf')
  m.connect()

  tasksq, resultsq, masterq, slavesq = get_qs(m)

  print "wait on all participants"
  for i in xrange(expected_participants):
    worker = masterq.get()
    print "%d of %d : %s is ready" % (i, expected_participants - 1, worker)

  print "everyone ready!"

  results = []

  submit_start = time.time()
  for i in xrange(count):
    X = np.random.random((1,3))[0]
    tasksq.put((i, options.order, options.extra, X))
  submit_end = time.time()

  for i in xrange(expected_participants):
    print "sending %d th start message" % i
    slavesq.put("start")

  receive_start = time.time()
  widgets = ['Progress: ', Percentage(), ' ', Bar(), ' ', ETA()]
  pbar = ProgressBar(widgets = widgets, maxval = count)
  pbar.start()
  for i in xrange(count):
    results.append(resultsq.get())
    pbar.update(i+1)
  receive_end = time.time()
  pbar.finish()

  submit  = submit_end - submit_start
  receive = receive_end - receive_start

  # shut down all participants
  for i in xrange(expected_participants):
    if options.last:
      slavesq.put("teardown")
    else:
      slavesq.put("prepare for more")

  # post processing
  stats = {}
  stats['submit'               ] = float(submit)
  stats['receive'              ] = float(receive)
  stats['count'                ] = count
  stats['expected_participants'] = expected_participants

  print "%s" % stats
  log.error("stats: %s", stats)

  tasks_accomplished_by = defaultdict(int)
  for i in results:
    tasks_accomplished_by[i[1]] += 1
  stats['tasks'] = tasks_accomplished_by

  npresults = np.array([(i[0],i[2],i[3],i[4], i[5]) for i in results])

  s = shelve.open(options.shelvename)
  n = str(time.time())
  s[n] = {
    'stats'   : stats,
    'results' : npresults,
  }
  s.close()