forward-selection

#!/usr/bin/env python

import collections
import heapq
import itertools
import json
import math
import numpy as np
import operator
import os
import pickle
import random
import redis
import scipy
from sklearn.linear_model import Ridge
import signal
import socket
import subprocess
import sys
import time
import uuid
import xml.etree.ElementTree as ET


FAILED_SCORE = 0.0
REUSE_STATELESS_PARAMS = True


class CommonEqualityMixin(object):
  def __eq__(self, other):
    return isinstance(other, self.__class__) and self.__dict__ == other.__dict__
  def __ne__(self, other):
    return not self.__eq__(other)


class Feature(CommonEqualityMixin):
  def __init__(self, stage, syntax_mode, label, templates=[]):
    ''' label is just a priveledged template that has high selectivitity and must be there '''
    if type(templates) is str:
      templates = [templates]
    assert type(label) is str
    assert type(syntax_mode) is str
    assert type(stage) is str
    assert type(templates) is list # of strinsgs
    self.stage = stage
    self.syntax_mode = syntax_mode
    self.label = label
    self.templates = templates
    self.templates.sort()

  def __str__(self):
    if self.stage == 'NullStage':
      return '*'.join([self.label] + self.templates)
    st = self.stage + '-' + self.syntax_mode
    return "*".join([st, self.label] + self.templates)

  def estimate_cardinality(self, template_info):
    c = 1
    for t in self.templates:
      c *= template_info.cardinality(self.stage, self.syntax_mode, self.label, t)
    assert type(c) is int or type(c) is long
    return c

  def propose_modification(self, template_info):
    def helper():
      # the higher the estimated cardinality, the greater the chance of
      # removing a template
      # e.g. cardinality = 100  -> p(shrink) = 0.500
      #      cardinality = 4    -> p(shrink) = 0.167
      #      cardinality = 400  -> p(shrink) = 0.667
      #      cardinality = 4000 -> p(shrink) = 0.863 (only if #templates > 1 though)
      card = pow(self.estimate_cardinality(template_info), 0.5)
      p_shrink = card / (card + 10.0)
      if len(self.templates) > 1 and random.random() < p_shrink:
        # remove a template
        tmpl = random.sample(self.templates, len(self.templates) - 1)
        return Feature(self.stage, self.syntax_mode, self.label, tmpl)
      else:
        # add a new template
        xs = [x for x in template_info.basic_templates() if x not in self.templates]
        if len(xs) > 0:
          return Feature(self.stage, self.syntax_mode, self.label, self.templates + [random.choice(xs)])
        else:
          raise Exception('cant extend this Feature any further')
    for i in range(500):
      f = helper()
      if f.estimate_cardinality(template_info) > 0:
        return f
    assert False, 'couldn\'t make a new feature with card>0'
    return None


class FeatureSet(object):
  def __init__(self, features=[], labels=None, syntax_mode=None, only_emit_diff=False, derived_from=None):
    '''
    If only_emit_diff=True, then calls to config_string() will only return the Features
    which appear in this feature set and NOT in derived_from. For the purpose of storing
    the full feature set for modifications, the features list will still contain Features
    from derived_from.
    '''
    self.features = features
    self.derived_from = derived_from
    self.__syntax_mode = syntax_mode
    self.__labels = labels
    self.only_emit_diff = only_emit_diff

    if derived_from and syntax_mode is None and derived_from.__syntax_mode:
      self.__syntax_mode = derived_from.__syntax_mode
    if derived_from and labels is None and derived_from.__labels:
      self.__labels = derived_from.__labels

  def __eq__(self, other):
    return isinstance(other, self.__class__) and self.features == other.features

  def __ne__(self, other):
    return not self.__eq__(other)

  def copy(self):
    return FeatureSet(self.features[:], self.derived_from)

  def restrict_to(self, stages):
    n = len(self.features)
    fs = []
    for f in self.features:
      if f.stage in stages:
        fs.append(f)
    self.features = fs
    print "[FeatureSet restrict_to] only including stages: %s, %d => %d" \
      % (str(stages), n, len(self.features))

  def trim(self, template_info):
    ''' Removes any features in this set which have cardinality 0 '''
    n = len(self.features)
    fs = []
    for f in self.features:
      if f.estimate_cardinality(template_info) > 0:
        fs.append(f)
      else:
        print '[FeatureSet trim] removing:', f
    self.features = fs
    print "[FeatureSet trim] %d => %d" % (n, len(self.features))

  def estimate_cardinality(self, template_info, stage=None):
    ''' Returns a dict with stage names as keys and int/long values '''
    cards = collections.defaultdict(lambda: 1)
    for f in self.features:
      cards[f.stage] = cards[f.stage] * f.estimate_cardinality(template_info)
    if stage:
      return cards[stage]
    else:
      return cards

  def stages(self):
    s = set()
    for f in self.features:
      s.add(f.stage)
    return list(s)

  def labels(self, stage=None):
    ''' If a stage is not provided, we union over stages '''
    labs = set()
    for f in self.features:
      if stage is None or f.stage == stage:
        labs.add(f.label)
    return list(labs)

  def __count_feats_for_stage(self, stage):
    c = 0
    for f in self.features:
      if f.stage == stage:
        c += 1
    return c

  def propose_modification(self, template_info, ignore_stage=False, only_add=False):
    '''
    Returns a new FeatureSet derived from this feature set.

    If ignore_stage=true, then this will produce Features with stage='NullStage',
    which should have corresponding rows in the TemplateInfo file (which say what
    templates/labels/syntax_modes can be used with 'NullStage').
    '''

    # a few options for mutation:
    # 1) add a Feature
    # 2) remove a Feature
    # 3) mutate a Feature

    debug = True

    if ignore_stage:
      stage = 'NullStage'
    else:
      rel_stages = set(template_info.stages()) & set(self.stages())
      if len(rel_stages) == 0:
        print 'self.stages', self.stages()
        print 'template_info.stages', template_info.stages()
        assert False
      stage = random.choice(list(rel_stages))
    relevant_feats = [f for f in self.features if f.stage == stage]
    card = 1
    #card = self.estimate_cardinality(template_info, stage)
    #l = 0.3   # interpolation between cardinality (1) and #features (0)
    #alpha = 0.2
    #beta = 500.0
    #gamma = 15.0
    #p_add_feature = l * beta / (beta + pow(card, alpha)) \
    #  + (1.0 - l) * (gamma / (gamma + len(self.features)))
    p_add_feature = 0.4

    p_remove_feature = 0.3 if len(relevant_feats) > 0 else 0.0
    p_mutate_feature = 0.3

    # re-normalize
    z = p_add_feature + p_remove_feature + p_mutate_feature
    p_add_feature /= z
    p_remove_feature /= z
    p_mutate_feature /= z

    if debug:
      print '[FeatureSet propose_modification] len(self.features) =', \
        len(self.features), 'stage =', stage, 'cardinality =', card, \
        'p_add_features =', p_add_feature, \
        'p_remove_feature =', p_remove_feature, \
        'p_mutate_feature =', p_mutate_feature

    r = random.random()
    if r < p_add_feature or only_add:
      # Make a new feature and add it to this set
      f = self.new_feature(template_info, stage)
      if debug:
        print '[FeatureSet propose_modification] adding:', f
      return FeatureSet(self.features + [f], derived_from=self)
    elif r >= p_add_feature and r < p_add_feature + p_remove_feature:
      # Remove a feature from this set
      f = random.choice(relevant_feats)
      if debug:
        print '[FeatureSet propose_modification] removing:', f
      i = self.features.index(f)
      return FeatureSet(self.features[:i] + self.features[i+1:], derived_from=self)
    else:
      # Modify and replace a feature that is already in this feature set
      try:
        i = random.randrange(len(self.features))
        f = self.features[i].propose_modification(template_info)
        if debug:
          print '[FeatureSet propose_modification] mutating:', self.features[i]
        return FeatureSet([f] + self.features[:i] + self.features[i+1:], derived_from=self)
      except:
        return self.propose_modification(template_info, ignore_stage)

  def new_feature(self, template_info, stage):
    ''' Creates a feature that is not already in this feature set '''
    # prefer adding features with small number of templates
    syntax_mode = self.syntax_mode()
    label_template, basic_template = random.choice(template_info.labels_for(stage, syntax_mode))
    f = Feature(stage, syntax_mode, label_template, basic_template)
    if f in self.features or f.estimate_cardinality(template_info) == 0:
      return self.new_feature(template_info, stage)
    else:
      print '[new_feature] stage =', stage
      print '[new_feature] syntax_mode =', syntax_mode
      print '[new_feature] label_template =', label_template
      print '[new_feature] basic_template =', basic_template
      return f

  def syntax_mode(self):
    if self.__syntax_mode:
      return self.__syntax_mode
    assert len(self.features) > 0
    return self.features[0].syntax_mode

  def __str__(self):
    return "<FeatureSet %s>" % (', '.join([str(x) for x in self.features]))

  def config_string(self):
    rel = self.features
    if self.only_emit_diff and self.derived_from:
      # Only keep features not in derived_from
      old = set()
      for f in self.derived_from.features:
        old.add(str(f))
      rel = []
      for f in self.features:
        if str(f) not in old:
          rel.append(f)
    return "+".join([str(x) for x in rel])


class Config(CommonEqualityMixin):
  ''' A full specification of a FN parsing model '''
  def __init__(self, feature_set, settings_dict, regularizer_dict, batch_size_dict, passes_dict, derived_from=None):
    self.feature_set = feature_set

    # Score of this configuration once evaluated (e.g. an F1 score)
    # This value is not set by this class
    self.score = None
    self.derived_from = derived_from

    self.regularizer_dict = regularizer_dict
    for k, v in regularizer_dict.iteritems():
      assert type(k) is str # Name of the stage
      assert type(v) is int # Natural log of the regularizer

    self.batch_size_dict = batch_size_dict
    for k, v in batch_size_dict.iteritems():
      assert type(k) is str # Name of the stage
      assert type(v) is int # Batch size

    self.passes_dict = passes_dict
    for k, v in passes_dict.iteritems():
      assert type(k) is str # Name of the stage
      assert type(v) is int # Number of training passes

    self.params = {}
    self.params['randomSeed'] = '9001'
    self.params.update(settings_dict)

  def stages(self):
    s1 = set(self.regularizer_dict.keys())
    s2 = set(self.features.stages())
    assert s1.issuperset(s2)
    return list(s1)

  def __str__(self):
    return "<Config features=%s regularizer=%s params=%s>" % \
      (str(self.feature_set), str(self.regularizer_dict), str(self.params))

  def copy(self):
    r = {}
    r.update(self.regularizer_dict)
    b = {}
    b.update(self.batch_size_dict)
    p = {}
    p.update(self.passes_dict)
    x = {}
    x.update(self.params)
    fs = self.feature_set.copy()
    return Config(fs, x, r, b, p)

# NOTE: This is only here because pickle doesn't work on instance methods...
# need to do the 'kingdom of nouns' thing and make an object to make pickle work.
class Mutator(object):
  def mutate(self, item):
    ''' return a list of items derived from the given item '''
    raise NotImplementedError()

class ConfigMutator(Mutator):
  def __init__(self, template_info, ignore_stages=False, emit_feature_diffs=False):
    '''
    If emit_feature_diffs=True, then this class will return mutated FeatureSets with
    fs.only_emit_diff=True, which means that anyone who calls fs.config_string() will
    only see the diffs between fs and fs.derived_from.
    NOTE: I put this here instead of in ConfigJobStarter because it doesn't really
    belong there.
    '''
    self.__template_info = template_info
    self.ignore_stages = ignore_stages
    self.emit_feature_diffs = emit_feature_diffs

  def mutate(self, config):
    ''' Config -> [Config] '''

    # TODO keep track of how many results we've seen so far
    # every K feature sets we get results for, the training set size should be increased by a factor of G
    # try K=50, G=3
    # this would take 5 iterations if N_0=1000 to get the full train set of 169k, 250 jobs (most of which would be quite big)
    # every time N*=G, consider lowering each stages regularizer (cross prod)
    # TODO does this mean that we need to block between hitting the limit (#feature-sets per N_train) and receiving all of the results?
    # could have a var that keeps track of the partition of jobs with the relevant N_train that we're submitting
    # still wouldn't fix the case of having the best (N-1)_train job occur at the end (we would have started N_train jobs already, which would have better perf... but wouldn't show up until later)
    # 1) hit num_jobs_at_this_train_size limit
    # 2) mutate_num_jobs_parent += 1
    # 3) start mutated jobs coming from *new* n_train limit pool

    # self.n_train_cur = only mutate configs that have n_train == this
    # self.n_jobs_dispatched = number of jobs submitted with the current setting self.n_train_cur
    # self.n_jobs_received = number of jobs submitted with the current setting self.n_train_cur

    ### OBSERVE_SCORE
    # if self.n_jobs_at_n_train_cur >= K:
    #   self.n_jobs_at_n_train_cur = 0
    #   self.n_train_cur *= G

    ### MUTATE
    # find the pool of jobs which match self.n_train_cur
    # choose one, mutate it, return it
    # self.n_jobs_at_n_train_cur += 1




    # CONSTRAINT:
    # We want to have had the chance to mutate every config in a given partition of n_train.
    # (or almost, we need to be failure resiliant)
    # if we assume that time(N_0 * G^{n+1}) >> time(N_0 * G^{n})
    
    # SOLUTION:
    # have a set of 'configs we havent tried mutating'
    # every time we attempt to mutate, we remove all the ones we considered
    # every time we receive a new score, we add it to this set
    # these set(s) are indexed by n_train
    # always try to select from smallest n_train bucket if its not empty
    # when one is not empty, we consider that example as well as any that we know about that have the same n_train

    # TODO this should not go here, it should go into a new Queue implementation
    # mutate just mutates, this new queue implementation will tell it what to mutate

    if self.ignore_stages:
      next_fs = config.feature_set.propose_modification(self.__template_info, ignore_stage=True)
    else:
      next_fs = config.feature_set.propose_modification(self.__template_info)

    if self.emit_feature_diffs:
      print '[mutate_config] only emitting FeatureSet diffs'
      next_fs.only_emit_diff = True

    cur_card = config.feature_set.estimate_cardinality(self.__template_info)
    next_card = next_fs.estimate_cardinality(self.__template_info)
    print '[mutate_config] cardinality: ', cur_card, ' => ', next_card
    # cardinality is now a dict with stages as keys
    # regularizer is too
    # for each stage, decide if you want to change the regularizer or not
    alpha = 15000
    beta = 1.5
    gamma = 0.25   # p(change regularizer irrespective of cardinality)
    children = [Config(next_fs, config.params, config.regularizer_dict, config.batch_size_dict, config.passes_dict, derived_from=config)]
    stages = set(self.__template_info.stages()) & set(next_fs.stages())
    for stage in stages:
      cc = cur_card[stage]
      nc = next_card[stage]
      dc = nc - cc
      try:
        if dc > alpha or (cc > 0 and dc / cc > beta) or random.random() < gamma:
          # stronger regularizer
          reg = {}
          reg.update(config.regularizer_dict)
          reg[stage] = reg[stage] + 1
          c = Config(next_fs, config.params, reg, config.batch_size_dict, \
            config.passes_dict, derived_from=config)
          children.append(c)
          print '[mutate] upping regularizer for', stage, 'to', reg[stage]
        if dc < -alpha or (cc > 0 and dc / cc < -beta) or random.random() < gamma:
          # weaker regularizer
          reg = {}
          reg.update(config.regularizer_dict)
          reg[stage] = reg[stage] - 1
          c = Config(next_fs, config.params, reg, config.batch_size_dict, \
            config.passes_dict, derived_from=config)
          children.append(c)
          print '[mutate] downing regularizer for', stage, 'to', reg[stage]
      except:
        print '[mutate] WARNING: goofed for stage', stage
        pass


    # Carry over stateless params learned at the previous iteration
    if REUSE_STATELESS_PARAMS:
      for c in children:
        if c.derived_from and 'workingDir' in c.derived_from.params:
          wd = c.derived_from.params['workingDir']
          f = os.path.join(wd, 'statelessParams.jser.gz')
          if not os.path.isfile(f):
            print 'look for', f, 'in', wd
            #raise Exception('templated stateless params not saved!')
            print 'WARNING: templated stateless params not saved!'
          else:
            c.params['fixedStatelessParamsJserFile'] = f
            print 'setting fixedStatelessParamsJserFile =', f
        else:
          print 'cant set fixedStatelessParamsJserFile because no derived_from'


    # TODO try increasing/decreasing the number of passes
    return children

class Queue(object):
  def observe_score(self, score, name, item):
    ''' Receive a message about the final score of an item running '''
    #raise NotImplementedError()
    pass
  def pop(self):
    raise NotImplementedError()

class ExplicitQueue(Queue):
  ''' Takes a list of items to pop in order '''
  def __init__(self, items):
    self.items = items
  def pop(self):
    if self.items:
      r = self.items[0]
      self.items = self.items[1:]
      return r
    else:
      return None

class MutatorQueue(Queue):
  '''
  An infinite queue of items. A seed set of items should be pushed onto the queue
  to begin with, but after that this queue will generate new items by calling
  propose_modification on existing items. Remember to call observe_score to let this
  queue know what are good items and which are bad.
  '''
  def __init__(self, mutator, greediness=1.0, max_pops=0):
    self.__mutator = mutator  # see Mutator.mutate, item -> [item]
    self.greediness = greediness
    self.item2score = {}    # item -> score for popped items
    self.waiting = []       # List of items that should be popped before mutating scored items
    self.__max_pops = max_pops
    self.__pop_count = 0
    self.stopped = False

  def _pick(self, greediness):
    debug = len(self.item2score) < 25
    # Sort the items by score, and find the max score
    weights = []
    items = []
    x = sorted(self.item2score.items(), key=lambda kv: kv[1], reverse=True)
    m = x[0][1]
    z = 0.0
    rank = 0
    for item, s in x:
      regret_a = (m - s) * 10.0
      regret_b = math.sqrt(rank) / 10.0
      w = math.exp(-greediness * (regret_a + regret_b))
      weights.append(w)
      items.append(item)
      z += w
      rank += 1
    t = random.random() * z
    if debug:
      print '[pick] z =', z, 't =', t
    # Randomly sample
    c = 0.0
    for i in range(len(weights)):
      c += weights[i]
      if debug:
        print '[pick]', items[i], 'has score', c
      if c >= t:
        print '[pick] chose', items[i]
        return items[i]
    # Should never get here
    print '[pick] scored_feature_sets:', self.item2score
    print '[pick] weights:', weights
    assert False

  def message(self, message):
    if message == 'stop':
      self.stopped = True
    else:
      print '[MutatorQueue message] unknown message:', message

  def show_state(self, k=10):
    print '[show_state]', len(self.item2score), 'observations'
    print '[show_state] stopped:', self.stopped
    best = sorted(self.item2score.items(), key=operator.itemgetter(1), reverse=True)
    i = 1
    for item, score in best[:k]:
      print "[show_state] %dth best\t%.3f\t%s" % (i, score, item)
      i += 1

  def observe_score(self, score, name, item):
    ''' Record the score of a given config/item '''
    print '[observe_score] score =', score, 'item =', item
    assert type(score) is float
    assert type(name) is str
    assert type(item) is not str  # TODO for debugging, in the future this could technically be a string

    try:
      item.score = score
    except:
      print '[observe_score] this item type does not have a score attribute:', type(item)

    if item in self.item2score:
      if self.item2score[item] == FAILED_SCORE:
        print ("[observe_score] WARN: may need to sleep longer, " \
          "%s/%s was originally deemed failed but later reported a score of %f") \
          % (item, name, score)
      else:
        raise Exception(str(item) + '/' + name + ' had a score of ' \
          + str(self.item2score[item]) \
          + ' but you tried to observe the score ' + str(score))
    self.item2score[item] = score
    # print the biggest weights
    if len(self.item2score) % 50 == 0:
      self.show_state()

  def seed(self, item):
    ''' Like push, but this item is special in that we will use it as the seed
    for all items until we start receiving scores back '''
    print '[seed] seeding', item, 'with score', FAILED_SCORE
    self.push(item)
    self.observe_score(FAILED_SCORE, 'seed' + str(len(self.item2score)), item)

  def push(self, item):
    ''' Push an item onto the queue which must be run before
    propose_modification's are considered '''
    self.waiting.append(item)

  def pop(self):
    ''' Returns a pushed item if there is one, otherwise chooses
    a scored item and mutates it with propose_modification '''
    if self.stopped:
      print '[pop] MutatorQueue is stopped, returning None'
      return None
    if self.__max_pops > 0 and self.__pop_count >= self.__max_pops:
      print '[pop]', str(self), 'hit max pops:', self.__max_pops
      return None
    self.__pop_count += 1
    if self.waiting:
      print '[pop] from waiting'
      return self.waiting.pop()
    else:
      if len(self.item2score) == 0:
        raise Exception('you must seed some items using push so there is something to mutate')
      tries = 0
      max_tries = 100
      while tries < max_tries:
        # g->0 as we run out of tries
        g = self.greediness * (max_tries - tries) / float(max_tries)
        parent = self._pick(g)
        children = self.__mutator.mutate(parent)
        dont_generate = set(self.waiting + self.item2score.keys())
        feasible = [c for c in children if c not in dont_generate]
        if len(feasible) < len(children):
          for c in (set(children) - set(feasible)):
            print '[pop] pruned', c
        if feasible:
          random.shuffle(feasible)
          while len(feasible) > 1:
            self.push(feasible.pop())
          print '[pop] after', tries, 'tries'
          return feasible.pop()
        else:
          tries += 1
      print '[pop] couldn\'t generate a new configuration'
      return None

class MultiQueue(object):
  ''' Muxes multiple Queues into one Queue '''
  def __init__(self):
    self.__item2qn = {}
    self.__name2q = {}
    self.__q_names = []
    self.__on_deck = 0
    self.__stopped = set()  # names of queues that are paused/stopped
    self.__all_stopped = False

  def add_queue(self, name, queue):
    assert name not in self.__name2q
    self.__name2q[name] = queue
    self.__q_names.append(name)

  def message(self, message):
    print '[MultiQueue message]', message
    if message == 'list' or message == 'info':
      if self.__all_stopped:
          print '[MultiQueue message list] all stopped:', self.__q_names
      elif self.__q_names:
        for q in self.__q_names:
          r = 'stopped' if q in self.__stopped else 'running'
          print '[MultiQueue message list]', q, 'is', r
      else:
        print '[MultiQueue message list] no queues!'
    elif message == 'stop':
      self.__all_stopped = True
    elif message == 'start':
      self.__all_stopped = False
    elif message == 'help':
      print '[MultiQueue message help] allowable commands:'
      print 'list: list all queues'
      print 'help: this command'
      print 'stop: pause all queues'
      print 'start: resume all queues'
      print '[queue_name] [command]: dispatch command to queue_name'
    else:
      # queue-specific message
      try:
        qn, msg = message.split(' ', 1)
        if msg == 'stop':
          print '[MultiQueue message] stopping', qn
          self.__stopped.add(qn)
        elif msg == 'start':
          print '[MultiQueue message] starting', qn
          self.__stopped.remove(qn)
        else:
          print '[MultiQueue message] forwarding', msg, 'to', qn
          q = self.__name2q[qn]
          q.message(msg)
      except:
        print '[MultiQueue message] failed to handle message:', message

  def observe_score(self, score, name, item):
    # parse name, lookup queue, forward call
    qn = self.__item2qn[item]
    q = self.__name2q[qn]
    q.observe_score(score, name, item)

  def __nextq(self):
    n = len(self.__q_names)
    if self.__on_deck == n - 1:
      self.__on_deck = 0
    else:
      self.__on_deck += 1

  def pop(self):
    if self.__all_stopped:
      print '[MultiQueue pop] stopped, returning None'
      return None
    n = len(self.__q_names)
    assert self.__on_deck < n
    for i in range(n):
      qn = self.__q_names[self.__on_deck]
      if qn in self.__stopped:
        self.__nextq()
      else:
        q = self.__name2q[qn]
        p = q.pop()
        self.__nextq()
        if p:
          self.__item2qn[p] = qn
          return p
    print '[MultiQueue pop] returning None', len(self.__stopped), 'stopped queues'
    return None

class SgeJobTracker(object):
  ''' A job tracker that asks qstat for the jobs that are
  running and spawns jobs with qsub '''
  def __init__(self, logging_dir=None):
    self.logging_dir = logging_dir
    if logging_dir and not os.path.isdir(logging_dir):
      os.makedirs(logging_dir)

  def can_submit_more_jobs(self):
    try:
      return len(self.jobs_queued()) < 10
    except subprocess.CalledProcessError:
      return False

  def jobs(self):
    '''
    name_predicate should be a lambda that takes a string (name)
    and returns true if the job should be kept.
    This method skips over any jobs that are marked as QLOGIN,
    so name_predicate need not filter those out.
    Returns a list of job names.
    '''
    xml = subprocess.check_output(['qstat', '-u', 'twolfe', '-xml'])
    xml = ET.fromstring(xml)
    assert xml.tag == 'job_info'
    # NOTE: wow this is really bad...
    # SGE reports *running* jobs in a list called 'queue_info'
    # and reports *queued* jobs in a list called 'job_info'
    for info_name in ['job_info', 'queue_info']:
      info = xml.find(info_name)
      assert info is not None
      # NOTE: each 'job_list' is actually a job
      # not a list of jobs as the name would suggest
      for j in info.findall('job_list'):
        #print 'j.tag', j.tag
        state = j.find('state').text    # e.g. 'r' or 'qw'
        name = j.find('JB_name').text
        #print '[sge jobs]', state, name
        if name == 'QLOGIN':
          continue
        yield (state, name)

  def jobs_running(self):
    ''' returns a list of job names '''
    return [name for state, name in self.jobs() if state == 'r']

  def jobs_queued(self):
    ''' returns a list of job names '''
    return [name for state, name in self.jobs() if state == 'qw']

  def spawn(self, name, args):
    cmd = ['qsub', '-N', name]
    if self.logging_dir:
      print 'spawning with logging dir:', self.logging_dir
      cmd += ['-o', self.logging_dir]
    # TODO add this as an option!
    #cmd += ['ForwardSelectionWorker.qsub', name]
    #cmd += ['FinalResults', '-1']
    cmd += ['scripts/propbank-feature-selection.sh', name]
    for k, v in args.iteritems():
      cmd += [k, v]
    print '[sge spawn] cmd =', cmd
    subprocess.Popen(cmd)
    time.sleep(0.2)


class LocalJobTracker(object):
  '''
  mock job tracker which uses redis instead of qsub
  if debug is true, this will call scripts/dummy-forward-selection-job
  else this will call the actual experiment
  '''
  def __init__(self, debug=False, max_concurrent_jobs=2):
    self.key = 'dummy-job-tracker.jobs'
    self.redis = redis.StrictRedis(host='localhost', port=6379, db=0)
    self.debug = debug
    self.max_concurrent_jobs = max_concurrent_jobs

  def remove_all_jobs(self):
    ''' ensures that there are no jobs running (for testing) '''
    self.redis.delete(self.key)

  def can_submit_more_jobs(self):
    return len(self.jobs_running()) < self.max_concurrent_jobs

  def jobs_running(self):
    # TODO this won't work in cases where jobs die!
    # qsub can handle this, but to remove from redis queue, we've been assuming things finish
    return self.redis.lrange(self.key, 0, -1)

  def set_job_done(self, name):
    print '[set_job_done] name=' + name
    self.redis.lrem(self.key, 0, name)

  def jobs_queued(self):
    return []

  # NOTE this implementation is a bit annoying because stdout from the spawned process
  # dumps into the main process's stdout, BUT this is only a problem locally. This appears
  # difficult to solve in python because there is no way to properly close the file that
  # you would route stdout to.
  def spawn(self, name, args):
    ''' args should be a dictionary of values to pass to grid.Runner '''
    self.redis.rpush(self.key, name)
    if self.debug:
      raise Exception('need to update args for dummy-forward-selection-job')
      subprocess.Popen(['scripts/dummy-forward-selection-job'] + args)
    else:
      cp = subprocess.check_output(['find', 'target/', '-iname', '*.jar']).strip()
      cp = ':'.join([x.strip() for x in cp.split('\n')])
      #cp = 'target/classes:' + cp
      cmd = ['java', '-Xmx4G', '-ea', '-server', '-cp', cp]
      #cmd += ['-XX:+UseG1GC', '-XX:G1ReservePercent=2', '-XX:ConcGCThreads=1', '-XX:ParallelGCThreads=1']
      cmd += ['-XX:+UseSerialGC']
      #cmd.append('edu.jhu.hlt.fnparse.experiment.grid.Runner')
      cmd.append('edu.jhu.hlt.fnparse.rl.rerank.RerankerTrainer')
      cmd.append(name)
      for k, v in args.iteritems():
        cmd += [k, v]
      print '[spawn] about to spawn:', ' '.join(["'" + x + "'" for x in cmd])
      subprocess.Popen(cmd)

# JobEngine: listens to messages, pops items from a Queue, names jobs, and asks the JobTracker to start the job
# JobTracker: reports the state of the computation resource (grid), w.r.t. a user (possibly including items from multiple Queues)
# JobStarter: given a (name, item, JobTracker) start a job
class ConfigJobStarter:
  ''' A job starter for items that are Configs '''
  def __init__(self, working_dir, redis_config):
    self.working_dir = working_dir
    self.redis_config = redis_config

  def start_job(self, name, config, job_tracker):
    ''' Just a function that adds arguments that are not Config-specific,
    but constitute more information than a general JobTracker should know '''
    assert isinstance(config, Config), "%s has type %s" % (config, type(config))
    cmd = {}
    wd = os.path.join(self.working_dir, 'wd-' + name)
    rf = os.path.join(wd, 'results.txt')
    cmd['resultReporter'] = "redis:%s,%s,%s\tfile:%s" % \
      (self.redis_config['host'], self.redis_config['channel'], self.redis_config['port'], rf)
    cmd['workingDir'] = wd
    config.params['workingDir'] = wd
    cmd['features'] = config.feature_set.config_string()
    for stage, regularizer in config.regularizer_dict.iteritems():
      cmd['regularizer.' + stage] = str(math.exp(regularizer))
    for stage, batch_size in config.batch_size_dict.iteritems():
      cmd['batchSize.' + stage] = str(batch_size)
    for stage, passes in config.passes_dict.iteritems():
      cmd['passes.' + stage] = str(passes)

    # pull out the score of the previous config, and tell the job not to save its
    # model if the score is lower than that.
    if config.derived_from and config.derived_from.score:
      cmd['scoreToBeat'] = str(config.derived_from.score)

    latent_stages = set()
    for f in config.feature_set.features:
      if f.syntax_mode == 'latent':
        latent_stages.add(f.stage)
    for s in latent_stages:
      key = 'bpIters.' + s
      if key not in cmd:
        cmd[key] = '5'

    cmd.update(config.params)
    job_tracker.spawn(name, cmd)

class JobEngine:
  def __init__(self, name, job_tracker, config_q, job_starter, redis_config, poll_interval=8.0):
    print '[JobEngine] attempting to use redis server at', redis_config
    if not config_q:
      raise Exception('you cant give a None config_q')
    self.name = name
    self.job_tracker = job_tracker    # talks to qsub
    self.config_q = config_q          # provides pop() and observe_score(score, name, item)
    self.job_starter = job_starter    # has start_job(name, item, job_tracker)
    self.redis_config = redis_config
    self.name2item = {}
    self.poll_interval = poll_interval
    self.dispatched = set()           # names of the jobs that (should be) running

  def parse_message(self, data):
    '''
    parses a message from the experiment over redis pubsub
    and returns a tuple of (config, score)
    '''
    toks = data.split('\t', 2)
    assert len(toks) == 3
    score = float(toks[0])
    name = toks[1]
    config = toks[2]
    return (score, name, config)

  def start_job(self):
    ''' returns a unique job name and updates self.name2item '''
    # TODO move this to a separate method that is passed in, like mutate
    item = self.config_q.pop()
    if item:
      # TODO this name needs to start with the queue name so MultiQueue knows
      # what to do when you call .observe_score()
      name = "fs-%s-%d" % (self.name, len(self.name2item))
      assert name not in self.name2item
      self.name2item[name] = item
      self.job_starter.start_job(name, item, self.job_tracker)
      self.dispatched.add(name)
      print '[start_job] name =', name
      print '[start_job] params =', item.params
      return name
    else:
      print '[start_job] there are no more jobs, returning None'
      return None

  def __handle_message(self, message, perf_file):
    ''' NOTE there is no command for 'stop'. This is because you should use 'messageQ',
    tell it to 'stop', and then this loop will exit when all the jobs are done. '''
    toks = message.split(' ', 1)
    if len(toks) == 1:
      print '[__handle_message] unknown message type:', message
      return
    m_type, rest = toks
    if m_type == 'result':
      # parse out (score, name, config), do what we were doing before
      score, name, config = rest.split('\t', 2)
      score = float(score)
      item = self.name2item[name]
      print '[__handle_message]', name, '/', config, 'finished successfully with a score', score
      perf_file.write("%f\t%s\t%s\n" % (score, name, config))
      perf_file.flush()
      self.config_q.observe_score(score, name, item)
      # Remove this jobs from dispatched
      print '[__handle_message] about to print job tracker'
      print '[__handle_message]', self.job_tracker
      if type(self.job_tracker) is LocalJobTracker: # for debugging
        print '[__handle_message] stopping', name, 'on local job tracker'
        self.job_tracker.set_job_done(name)
      try:
        self.dispatched.remove(name)
        print '[__handle_message] removed', name, 'from dispatched,', len(self.dispatched), 'still dispatched'
      except:
        print '[__handle_message]', name, 'was not in dispatched, we gave up on this job as failed previously'
    elif m_type == 'messageQ':
      # check for a queue name argument, forward to q.message(remaining_message)
      try:
        self.config_q.message(rest)
      except (AttributeError, TypeError) as e:
        print '[__handle_message] this queue does not support messages:', rest
        print e
    elif m_type == 'saveQ':
      # parse out file, pickle the queue to that
      print '[__handle_message] saving queue to', rest
      with open(rest, 'wb') as f:
        pickle.dump(self.config_q, f)
    elif m_type == 'loadQ':
      # parse out file, unpickle the queue from that
      print '[__handle_message] loading queue from', rest
      with open(rest, 'rb') as f:
        self.config_q = pickle.load(f)
      self.config_q.show_state()
    else:
      print '[__handle_message] unknown message:', message

  def run(self, perf_file_name):
    r = redis.StrictRedis(host=self.redis_config['host'], port=self.redis_config['port'], db=self.redis_config['db'])
    p = r.pubsub(ignore_subscribe_messages=True)
    p.subscribe(self.redis_config['channel'])
    print '[run] writing results to', perf_file_name
    perf_file = open(perf_file_name, 'w')
    self.num_jobs = 0
    while True:
      # Try to dispatch new jobs
      can_sub = self.job_tracker.can_submit_more_jobs()
      if can_sub and self.start_job():
        self.num_jobs += 1
      else:
        # Check for results 
        if can_sub:
          print '[run] can submit more jobs, but queue is empty'
        else:
          print '[run] can\'t submit any more jobs'
        print '[run] checking on current jobs'
        message = p.get_message()
        if message:
          print '[run] received message:', message
          self.__handle_message(message['data'], perf_file)
        else:
          # check if any jobs died
          time.sleep(self.poll_interval)
          print '[run] no one phoned home, maybe someone died?'
          try:
            r = set(self.job_tracker.jobs_running() + self.job_tracker.jobs_queued())
            if len(r) == 0 and can_sub:
              print '[run] all jobs are done!'
              break
            print '[run] running and queued:', r
            failed = set([name for name in self.dispatched if name not in r])
            for name in failed:
              item = self.name2item[name]
              print '[run] dead:', name, item
              self.config_q.observe_score(FAILED_SCORE, name, item)
              self.dispatched.remove(name)
            else:
              print '[run] everything is running nicely:', len(r)
          except subprocess.CalledProcessError:
            print '[run] qstat failed! ignoring...'
    perf_file.close()


class TemplateInfo(object):
  # TODO could support cardinalities on higher-order polynomials of features
  # (e.g. products of templates), just add more columns
  def __init__(self, filename):
    self.filename = filename
    self.__info = {}  # (stageName, syntaxMode, labelTemplate, basicTemplate) -> (cardinaityInt, timeFloat)
    self.__templates = set()
    self.__stages = set()
    with open(filename, 'r') as f:
      for line in f:
        toks = line.strip().split('\t')
        assert len(toks) == 6, 'toks=' + str(toks)
        #template, stage, syntax_mode, card, time = toks
        stage, syntax_mode, label, template, card, time = toks
        k = (stage, syntax_mode, label, template)
        card = int(card)
        time = float(time)
        if k in self.__info:
          #raise Exception('duplicate key: ' + str(k))
          print 'WARNING: duplicate:', k, '\t(card, time) old:', self.__info[k], 'new:', (card, time)
          if card > self.__info[k][0]:
            self.__info[k] = (card, time)
        else:
          self.__info[k] = (card, time)
          self.__templates.add(template)
          self.__stages.add(stage)
        #print template, stage, card, time
    self.__templates = list(self.__templates)
    self.__stages = list(self.__stages)
    if len(self.__stages) == 0:
      raise Exception('no stages?')

    # compute the set features that have non-zero cardinality (and we might want to instantiate)
    self.__nonzero = collections.defaultdict(list)
    for (stage, syntax_mode, label, template), (card, time) in self.__info.iteritems():
      k = (stage, syntax_mode)
      v = (label, template)
      if card > 0:
        self.__nonzero[k].append(v)
      #else:
      #  print template, 'has no support in', k

    # stage -> [label/template]
    # if a stage appears here, only allow labels/templates to be chosen from this list
    #self.__label_restrictions = {}

  def cardinality(self, stage, syntax_mode, label_template, basic_template):
    k = (stage, syntax_mode, label_template, basic_template)
    # For debugging, we should get a KeyError if we can't lookup a template,
    # but I don't want this to halt the entire server if we're running experiments
    # "in production".
    #return self.__info[k][0]
    try:
      return self.__info[k][0]
    except:
      print '[cardinality] WARNING key error for:', k
      return 2

  def stages(self):
    return self.__stages

  def restrict_stages_to(self, stages):
    self.__stages = list(set(self.__stages) & set(stages))
    if len(self.__stages) == 0:
      raise Exception('no stages left after intersecting with: ' + str(stages))

  def labels_for(self, stage, syntax_mode):
    ''' Returns a list of (label_template, basic_template) which can be used with the given stage '''
    if stage == 'NullStage':
      # Take all templates that at least match the syntax mode
      nz = []
      for (stage_i, syntax_i), list_nz in self.__nonzero.iteritems():
        #print 'stage_i', stage_i, 'syntax_i', syntax_i, 'stage', stage
        if syntax_i == syntax_mode:
          nz += list_nz
    else:
      k = (stage, syntax_mode)
      print '[labels_for]', k
      print '[labels_for] WARNING: ignoring restrict_stages_to!'
      nz = self.__nonzero[k]
    if len(nz) == 0:
      print 'WARNING: found no labels for:', (stage, syntax_mode)
    return nz

  def basic_templates(self):
    return self.__templates

def generate_all_features(stages, syntax_mode, \
    labels=['frameRoleArg', 'roleArg', 'arg', 'framePrune', 'prune']):

  assert type(stages) is list
  assert type(syntax_mode) is str
  fs = []

  for fra in ['frameRoleArg', 'roleArg', 'arg', 'framePrune', 'prune']:
    if fra not in labels:
      continue
    fs.append((fra, ['1']))
    fs.append((fra, ['head1Word']))
    fs.append((fra, ['head1Pos2']))
    fs.append((fra, ['head1WordWnSynset']))
    fs.append((fra, ['head1Lemma']))
    fs.append((fra, ['span1span2Overlap']))
    fs.append((fra, ['span1LeftPos2']))
    fs.append((fra, ['span1FirstPos2']))
    fs.append((fra, ['span1LastPos2']))
    fs.append((fra, ['span1RightPos2']))
    fs.append((fra, ['span1Width/3']))
    fs.append((fra, ['span1Width/5']))
    fs.append((fra, ['Dist(SemaforPathLengths,Head1,Head2)']))
    fs.append((fra, ['head1CollapsedLabel']))
    fs.append((fra, ['head1CollapsedParentDir', 'Dist(Direction,Head1,Head2)']))
    fs.append((fra, ['head1head2Path-LEMMA-DIRECTION-t']))
    fs.append((fra, ['head1head2Path-POS-DEP-t']))
    fs.append((fra, ['span1StanfordCategory']))
    fs.append((fra, ['span1StanfordRule']))
    fs.append((fra, ['CfgFeat-CommonParent-Category']))
    fs.append((fra, ['CfgFeat-CommonParent-DeltaDepth']))

  if 'span1IsConstituent' in labels:
    fs.append(('span1IsConstituent', ['1']))
    fs.append(('span1IsConstituent', ['frameRoleArg']))
    fs.append(('span1IsConstituent', ['roleArg']))
    fs.append(('span1IsConstituent', ['frame']))
    fs.append(('span1IsConstituent', ['frame', 'Dist(Direction,Head1,Head2)']))
    fs.append(('span1IsConstituent', ['span1PosPat-COARSE_POS-1-1']))
    fs.append(('span1IsConstituent', ['span1PosPat-WORD_SHAPE-1-1']))
    cparts = ['span1LeftPos2', 'span1FirstPos2', 'span1LastPos2', 'span1RightPos2', 'span1Width/5', 'span1FirstWord', 'span1LastWord']
    for p1 in cparts:
      fs.append(('span1IsConstituent', [p1]))
      for p2 in cparts:
        if p2 <= p1:
          continue
        if 'Word' in p1 and 'Word' in p2:
          continue;
        fs.append(('span1IsConstituent', [p1, p2]))
        for p3 in cparts:
          if p3 <= p2:
            continue
          if 'Word' in p3 and ('Word' in p1 or 'Word' in p2):
            continue
          fs.append(('span1IsConstituent', [p1, p2, p3]))

  # NOTE: there are already unary dep link factors from pacaya not included here
  if 'head2GovHead1' in labels:
    fs.append(('head2GovHead1', ['1']))
    fs.append(('head2GovHead1', ['frameRoleArg']))
    fs.append(('head2GovHead1', ['frame', 'Dist(Direction,Head1,Head2)']))
    fs.append(('head2GovHead1', ['frame', 'head1Pos']))

  for s in stages:
    for label, templates in fs:
      yield Feature(s, syntax_mode, label, templates)

def seeds(parser_mode, syntax_mode, settings, template_info):
  ''' Generator of Configs '''
  p = 5  # default number of training passes (this value is copied for all runs)
  if parser_mode == 'span':
    s1 = 'RoleSpanPruningStage'
    s2 = 'RoleSpanLabelingStage'
  elif parser_mode == 'head':
    s1 = 'RoleHeadStage'
    s2 = 'RoleHeadToSpanStage'
  else:
    raise Exception('unknown parser_mode: ' + parser_mode)
  batch_size = {s1: 1, s2: 1}
  passes = {s1: p, s2: p}
  fs = []
  for f in generate_all_features([s1, s2], syntax_mode):
    if f.estimate_cardinality(template_info) > 0:
      fs.append(f)
  print "[seeds] seed set for %s-%s has %d features covering the following labels: %s" \
    % (parser_mode, syntax_mode, len(fs), str(list(set([f.label for f in fs]))))

  for reg1 in [11, 13]:
    for reg2 in [11, 13]:
      reg = {s1: reg1, s2: reg2}
      yield Config(FeatureSet(fs), settings, reg, batch_size, passes)

      # make a few more seeds which are random subsets of those features
      n = 5
      for p_keep in [0.5, 0.2]:
        for i in range(n):
          fs_smaller = fs[:]
          random.shuffle(fs_smaller)
          k = int(math.ceil(p_keep * len(fs)))
          fs_smaller = fs_smaller[:k]
          yield Config(FeatureSet(fs_smaller), settings, reg, batch_size, passes)

def experiments(parser_mode, syntax_mode, template_info_filename, max_train=99999, extra_settings=None):
  ''' Returns a (TemplateInfo, [Config]) '''
  assert parser_mode in ['span', 'head']
  assert syntax_mode in ['regular', 'latent', 'none']
  stages = {
    'span': ['RoleSpanPruningStage', 'RoleSpanLabelingStage'],
    'head': ['RoleHeadStage', 'RoleHeadToSpanStage']
  }
  template_info = TemplateInfo(template_info_filename)
  settings = {
    'parserMode': parser_mode,
    'syntaxMode': syntax_mode,
    'evaluationFunction': 'ArgumentMicroF1',
    'MaxTrainSize': str(max_train),
    'KpTrainDev': '1\t0.3',
  }
  if extra_settings:
    settings.update(extra_settings)
  s1, s2 = stages[parser_mode]
  sd = list(seeds(parser_mode, syntax_mode, settings, template_info))

  # Remove any irrelevant features that are in the seeds
  for conf in sd:
    print '[experiments] trimming feature set for', conf.params
    conf.feature_set.trim(template_info)

  return (template_info, sd)

def default_config(parser_mode, syntax_mode, template_info_filename):
  ''' Returns a Config for a default/good configuration of the model specified by params '''
  template_info, configs = experiments(parser_mode, syntax_mode, template_info_filename)
  return configs[0]

def propbank_feature_selection(template_info_filename, local):
  assert isinstance(local, bool)
  mq = MultiQueue()
  template_info = TemplateInfo(template_info_filename)

  # I believe NoneStage is used so that we never mutate the stage when we mutate the features (travis,aug26)
  labels = ['frameRoleArg', 'roleArg']
  fs = generate_all_features(['NoneStage'], 'NoneSyntaxMode', labels=labels)
  for f in fs:
    print f
  #print 'fs:', ', '.join(str(x) for x in fs)
  fs = FeatureSet(features=list(fs), labels=labels, syntax_mode='regular')

  params = {
    'nTrain': 10000,
    'shard': 0, 'numShards': 5,   # prevents reading ALL of the data in, make sure that nTrain < N / numShards
    #'workingDir': '/tmp/foo',    # this is set by tge
    'addStanfordParses': False,
    'realTestSet': False,
    'propbank': True,
    'l2Penalty': 1e-8,
    'secsBetweenShowingWeights': 180,
    'trainTimeLimit': 60, #240,
    'featCoversFrames': False,
    'useGlobalFeatures': False,
    'data.wordnet': 'toydata/wordnet/dict',
    'data.embeddings': 'data/embeddings',
    'data.ontonotes5': 'data/ontonotes-release-5.0/LDC2013T19/data/files/data/english/annotations',
    'data.propbank.conll': '../conll-formatted-ontonotes-5.0/conll-formatted-ontonotes-5.0/data',
    'data.propbank.frames': 'data/ontonotes-release-5.0-fixed-frames/frames',
    'disallowConcreteStanford': 'false',    # apparently I haven't parsed everything yet...
    'redis.port.propbankParses': '6379',
    'redis.db.propbankParses': '0',
  }
  if local:
    params.update({
      'redis.host.propbankParses': 'localhost',
      'laptop': True,
    })
  else:
    params.update({
      'redis.host.propbankParses': 'r5n39',  # TODO take as arg!
      'laptop': False,
    })

  # Debugging
  #params['bailOutOfTrainingASAP'] = True

  # Convert value to strings
  params = {a:str(b) for a,b in params.iteritems()}

  regularizer = {}
  batch_size = {}
  passes = {}
  conf = Config(fs, params, regularizer, batch_size, passes)

  max_fs = 5000   # number of feature sets to try
  mutator = ConfigMutator(template_info, ignore_stages=True)

  q = MutatorQueue(mutator, greediness=1.0, max_pops=max_fs)
  q.seed(conf)

  mq.add_queue('propbank_fs', q)

  # Debugging
  if local:
    print fs
    sys.exit(1)

  return mq


def forward_selection_queue(template_info_filename):
  ''' Returns a MultiQueue for forward (feature) selection on each model '''
  #n_train = 2500
  mq = MultiQueue()
  #for parser_mode in ['span', 'head']:
  for parser_mode in ['span']:
    for use_cky_decoder in ['true', 'false']:
      for one_stage in ['true', 'false']:
        for syntax_mode in ['regular', 'latent', 'none']:
          if one_stage == 'true' and syntax_mode == 'regular':
            continue
          for n_train in [1000, 3000]:
            q_name = "%s-%s-%s-%s-%d" % (parser_mode, syntax_mode, one_stage, use_cky_decoder, n_train)
            if not (parser_mode == 'span' and syntax_mode == 'latent') and use_cky_decoder != 'false':
              continue
            template_info, configs = experiments(parser_mode, syntax_mode, \
              template_info_filename, max_train=n_train, \
              extra_settings={ \
                'useCkyDecoder': use_cky_decoder, \
                'oneStage': one_stage})
            print 'setting up queue for', q_name
            print template_info
            print '\n\t'.join(str(c) for c in configs)

            # Create the config queue
            max_fs = 5000   # number of feature sets to try
            mutator = ConfigMutator(template_info)
            q = MutatorQueue(mutator, greediness=1.0, max_pops=max_fs)
            assert len(configs) > 0
            for conf in configs:
              q.seed(conf)
            mq.add_queue(q_name, q)
  return mq

def sweep(starting_config, dimensions_to_sweep):
  ''' dimensions_to_sweep should be a list of ('param', ['val1', 'val2', ...])
  where the semantics is that jobs will be launched as if dimension at the
  start of the loop will be treated as the outermost loop.
  You can also provide a dict, which will be converted to the first format
  in an arbitrary order. '''
  if isinstance(dimensions_to_sweep, dict):
    return sweep(starting_config, dimensions_to_sweep.items())
  assert isinstance(dimensions_to_sweep, list)
  jobs = []
  def pair(key, values):  # {k: [x,y,x], ...} => [[(k,x), (k,y), (k,z)], ...]
    return [(key, v) for v in values]
  flat = [pair(k,v) for k,v in dimensions_to_sweep]
  for conf in itertools.product(*flat):
    c = starting_config.copy()
    for k, v in conf:
      c.params[k] = v
    jobs.append(c)
  return ExplicitQueue(jobs)

def ab_test_queue(template_info_filename):
  # NOTE: These only need to be tested for latent.
  # under regular syntax vars are not rolled out, no binary factors
  mq = MultiQueue()

  # TODO Not doing this now because even if I kick this constraint up to
  # a ternary factor (arg * l_{target,arg} * l_{arg,target}), many args
  # are still not local (57% for *collapsed* dependencies).
  # Also, this stage appears to work pretty well already.
  #mq.add_queue('head-latent-s1-argIsCons',
  #  sweep(default_config('head', 'latent', template_info_filename),
  #    {'disallowArgWithoutDepenency.RoleHeadStage': ['true', 'false'],
  #     'MaxTrainSize': ['250', '750', '2250']}))

  # NOTE: MaxTrainSize should always come first so that we can sweep
  # the parameters over smaller datasets to get some initial results
  # before wasting a lot of time on the full thing.

  # IMPLMENTED IN JAVA
  # VISUALLY TESTED
  mq.add_queue('head-latent-s2-argIsCons',
    sweep(default_config('head', 'latent', template_info_filename),
      [('MaxTrainSize', ['500', '1500', '99999']),
       ('disallowArgWithoutConstituent.RoleHeadToSpanStage', ['true', 'false'])]))

  # IMPLMENTED IN JAVA
  # VISUALLY TESTED
  mq.add_queue('span-latent-s1-argIsCons',
    sweep(default_config('span', 'latent', template_info_filename),
      [('MaxTrainSize', ['500', '1500', '99999']),
       ('disallowArgWithoutConstituent.RoleSpanPruningStage', ['true', 'false'])]))

  # IMPLMENTED IN JAVA
  # VISUALLY TESTED
  mq.add_queue('span-latent-s2-argIsCons',
    sweep(default_config('span', 'latent', template_info_filename),
      [('MaxTrainSize', ['500', '1500', '99999']),
       ('disallowArgWithoutConstituent.RoleSpanLabelingStage', ['true', 'false'])]))

  # IMPLMENTED IN JAVA
  # VISUALLY TESTED
  mq.add_queue('span-regular-s1-pruneMethod',
    sweep(default_config('span', 'regular', template_info_filename),
      [('MaxTrainSize', ['500', '1500', '99999']),
       ('deterministicRolePruningMethod',
        ['STANFORD_CONSTITUENTS', 'DEPENDENCY_SPANS', 'XUE_PALMER',
          'XUE_PALMER_HERMANN', 'XUE_PALMER_DEP', 'XUE_PALMER_DEP_HERMANN'])]))

  mq.add_queue('span-bpIters',
    sweep(default_config('span', 'latent', template_info_filename),
      [('MaxTrainSize', ['500', '1500', '99999']),
       ('bpIters', ['1', '2', '4', '8', '16'])]))

  mq.add_queue('head-bpIters',
    sweep(default_config('head', 'latent', template_info_filename),
      [('MaxTrainSize', ['500', '1500', '99999']),
       ('bpIters', ['1', '2', '4', '8', '16'])]))

  return mq


def agiga_model_queue_frameid(template_info_filename):
  return agiga_model_queue(template_info_filename, True)

def agiga_model_queue_full(template_info_filename):
  return agiga_model_queue(template_info_filename, False)

def agiga_model_queue(template_info_filename, run_feature_selection=False):
  template_info, configs = experiments('span', 'regular', template_info_filename)
  c = configs[0]
  c.params['learnFrameId'] = 'true'
  c.params['MaxTrainSize'] = '99999'
  c.regularizer_dict['FrameIdStage'] = 12
  fs = [
    Feature('FrameIdStage', 'regular', 'frame', ['head1Word']),
    Feature('FrameIdStage', 'regular', 'frame', ['head1Bc1000/6']),
    Feature('FrameIdStage', 'regular', 'frame', ['head1CollapsedLabel']),
    Feature('FrameIdStage', 'regular', 'frame', ['head1Pos']),
    Feature('FrameIdStage', 'regular', 'frame', ['head1ParentPos']),
    Feature('FrameIdStage', 'regular', 'frame', ['head1ParentWord4']),
    Feature('FrameIdStage', 'regular', 'frame', ['head1RootPath-LEMMA-DEP-t']),
    Feature('FrameIdStage', 'regular', 'frame', ['Word4-2-grams-between-</S>-and-<S>']),
    Feature('FrameIdStage', 'regular', 'frame', ['span1LeftPos']),
    Feature('FrameIdStage', 'regular', 'frame', ['span1RightPos']),
    Feature('FrameIdStage', 'regular', 'frame', ['span1PosPat-COARSE_POS-3-0']),
    Feature('FrameIdStage', 'regular', 'frame', ['span1PosPat-COARSE_POS-0-3']),
  ]
  c.feature_set.features += fs
  if run_feature_selection:
    # Train only frameId
    c.params['evaluationFunction'] = 'TargetMicroF1'
    c.params['skipArgId'] = 'true'
    # prune any features that are not related to frameId
    c.feature_set.restrict_to(['FrameIdStage'])
    mut = ConfigMutator(template_info)
    q = MutatorQueue(mut)
    q.seed(c)
    return q
  else:
    # Train frameId + argLabeling
    #return ExplicitQueue([c])
    mut = ConfigMutator(template_info)
    q = MutatorQueue(mut)
    q.seed(c)
    return q

def debug_queue(template_info_filename):
  ''' a queue for whatever issue i'm currently debugging '''
  n = 1500
  template_info, configs = experiments('span', 'regular', template_info_filename, max_train=n)
  c_reg = configs[0]
  c_dup = c_reg.copy()
  f = Feature('RoleSpanLabelingStage', 'regular', 'frameRoleArg', ['span1StanfordCategory'])
  c_dup.feature_set.features.append(f)
  # NOTE this only works as long as I leave the bug in:
  # I had a bug in disallowArgWithoutConstituent where i enforced this when
  # there is no contituency tree and all spans are trivially "not constituents"
  c_reg_fix = c_reg.copy()
  c_reg_fix.params['disallowArgWithoutConstituent.RoleSpanLabelingStage'] = 'false'
  c_dup_fix = c_dup.copy()
  c_dup_fix.params['disallowArgWithoutConstituent.RoleSpanLabelingStage'] = 'false'
  return ExplicitQueue([c_reg_fix, c_reg, c_dup_fix, c_dup])


def run_experiments(template_info_filename, queue_name, working_dir, where):

  local = where == 'local'
  assert local or where == 'grid'

  if not os.path.isdir(working_dir):
    raise Exception('invalid working_dir: ' + working_dir)

  q_func = globals()[queue_name]
  if not q_func:
    raise Exception('unknown queue_name: ' + queue_name)
  print 'running', queue_name
  mq = q_func(template_info_filename, local)
  ## Work to be done
  #if job_name == 'fs-role':
  #  mq = forward_selection_queue(template_info_filename)
  #  working_dir = 'experiments/forward-selection/role/omni'
  #elif job_name == 'ab-test':
  #  mq = ab_test_queue(template_info_filename)
  #  working_dir = 'experiments/ab-test'
  #elif job_name == 'agiga':
  #  mq = agiga_model_queue(template_info_filename, run_feature_selection=False)
  #  working_dir = 'experiments/agiga'
  #elif job_name == 'debug':
  #  mq = debug_queue(template_info_filename)
  #  working_dir = 'experiments/debug'
  #else:
  #  raise Exception('unknown job_name: ' + job_name)
  if not mq:
    raise Exception('something went wrong retrieving: ' + queue_name)

  # Create the job tracker
  if local:
    job_tracker = LocalJobTracker(debug=False, max_concurrent_jobs=2)
    job_tracker.remove_all_jobs()
  else:
    d = os.path.join(working_dir, 'sge-logs')
    job_tracker = SgeJobTracker(logging_dir=d)

  perf_filename = os.path.join(working_dir, 'perf.txt')
  this_machine = socket.gethostname()
  redis_config = {
    'channel': 'forward-selection',
    'host': this_machine,
    'port': '6379',
    'db': '0'
  }
  job_starter = ConfigJobStarter(working_dir, redis_config)
  print 'working directory:', working_dir

  fs = JobEngine('test', job_tracker, mq, job_starter, redis_config)

  print 'starting...'
  fs.run(perf_filename)

if __name__ == '__main__':

  # Make stdout flush regularly
  sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)

  if len(sys.argv) == 5:
    run_experiments(*(sys.argv[1:]))
  else:
    print 'please provide:'
    print '1) a path to a basic template cardinality file (e.g. experiments/feature-cardinality-estimation/template.real.txt)'
    print '2) the name of a function that returns a queue, e.g. "propbank_feature_selection", "forward_selection_queue", or "agiga_model_queue"'
    print '3) a working directory'
    print '4) where to run (i.e. \"grid\" or \"local\")'
    sys.exit(-1)