#!/usr/bin/env python

import sys
import os
import time
from optparse import OptionParser

import numpy as np

from interp.grid.gmsh import ggrid
from interp.tools import (baker_exact_2D, baker_exact_3D,
                            friendly_exact_2D, friendly_exact_3D)
from interp.cluster import QueueManager, get_qs

if __name__ == '__main__':
    parser = OptionParser(usage="usage: %s [options] <server>"
                            " <gmsh file> <dimension> <baker|friendly>")

    parser.add_option("-v", "--verbose",
        action="store_true", dest="verbose", default=False,
        help="verbose flag (default: %default)")

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

    (options, args) = parser.parse_args()

    if len(args) != 4:
        parser.print_usage()
        sys.exit(1)

    server, input_file, dimension, exact_type = args
    dimension = int(dimension)

    if     dimension == 2 and exact_type == 'baker':
        exact = baker_exact_2D
    elif dimension == 2 and exact_type == 'friendly':
        exact = friendly_exact_2D
    elif dimension == 3 and exact_type == 'baker':
        exact = baker_exact_3D
    elif dimension == 3 and exact_type == 'friendly':
        exact = friendly_exact_3D

    myname = "%s-%d" % (os.uname()[1], os.getpid())
    if options.verbose:
        print "%s: started" % myname

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

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

    if options.verbose:
        print "%s: starting parse input file" % myname
    g = ggrid(input_file, dimension)
    g.q = np.array([exact(x) for x in g.verts])
    if options.verbose:
        print "%s: done parsing input file" % myname

    while True:
        if options.verbose:
            print "%s: letting master know that I am ready" % myname
        masterq.put(myname)

        if options.verbose:
            print "%s: waiting for master to tell me to start" % myname
        action = minionsq.get()
        if options.verbose:
            print "%s: got message from master: %s" % (myname, action)

        if action in ('teardown', 'slay'):
            masterq.get()
            break

        while not tasksq.empty():
            i, order, extra, X = tasksq.get()
            try:
                s = time.time()
                a = g.run_baker(X, order=order, extra_points=extra)
                cur_qlin = a.qlin
                cur_error = a.error
                cur_final = a.final
                cur_exact = exact(X)
                e = time.time()
                duration = e - s
            except Exception as e:
                print >>sys.stderr, X, e
                cur_qlin = 0.0
                cur_error = 0.0
                cur_final = 0.0
                cur_exact = 0.0
                duration = 0

            resultsq.put((i,
                    cur_qlin,
                    cur_error,
                    cur_final,
                    cur_exact,
                    duration,
                ))

    if options.verbose:
        print "%s: exiting" % myname