def __init__(self, expr,
                 name=None,
                 loss_target=None,
                 exceptions=None,
                 do_checks=True,
                 ):
        if do_checks:
            if isinstance(expr, pyll.Apply):
                self.expr = expr
                # XXX: verify that expr is a dictionary with the right keys,
                #      then refactor the code below
            elif isinstance(expr, dict):
                if 'loss' not in expr:
                    raise ValueError('expr must define a loss')
                if 'status' not in expr:
                    expr['status'] = STATUS_OK
                self.expr = pyll.as_apply(expr)
            else:
                raise TypeError('expr must be a dictionary')
        else:
            self.expr = pyll.as_apply(expr)

        self.params = {}
        for node in pyll.dfs(self.expr):
            if node.name == 'hyperopt_param':
                label = node.arg['label'].obj
                if label in self.params:
                    raise DuplicateLabel(label)
                self.params[label] = node.arg['obj']

        if exceptions is not None:
            self.exceptions = exceptions
        self.loss_target = loss_target
        self.name = name

    def __init__(self, expr, name=None, rseed=None, loss_target=None, exceptions=None, do_checks=True):

        if do_checks:
            if isinstance(expr, pyll.Apply):
                self.expr = expr
                # XXX: verify that expr is a dictionary with the right keys,
                #      then refactor the code below
            elif isinstance(expr, dict):
                if "loss" not in expr:
                    raise ValueError("expr must define a loss")
                if "status" not in expr:
                    expr["status"] = STATUS_OK
                self.expr = pyll.as_apply(expr)
            else:
                raise TypeError("expr must be a dictionary")
        else:
            self.expr = pyll.as_apply(expr)

        self.params = {}
        for node in pyll.dfs(self.expr):
            if node.name == "hyperopt_param":
                self.params[node.arg["label"].obj] = node.arg["obj"]

        if exceptions is not None:
            self.exceptions = exceptions
        self.loss_target = loss_target
        self.installed_rng = False
        if rseed is None:
            self.rng = None
        else:
            self.rng = np.random.RandomState(rseed)

        self.name = name

def test_vectorize_multipath():
    N = as_apply(15)

    p0 = hp_uniform('p0', 0, 1)
    loss = hp_choice('p1', [1, p0, -p0]) ** 2
    expr_idxs = scope.range(N)
    vh = VectorizeHelper(loss, expr_idxs, build=True)

    vloss = vh.v_expr
    print vloss

    full_output = as_apply([vloss,
        vh.idxs_by_label(),
        vh.vals_by_label()])

    new_vc = recursive_set_rng_kwarg(
            full_output,
            as_apply(np.random.RandomState(1)),
            )

    losses, idxs, vals = rec_eval(new_vc)
    print 'losses', losses
    print 'idxs p0', idxs['p0']
    print 'vals p0', vals['p0']
    print 'idxs p1', idxs['p1']
    print 'vals p1', vals['p1']
    p0dct = dict(zip(idxs['p0'], vals['p0']))
    p1dct = dict(zip(idxs['p1'], vals['p1']))
    for ii, li in enumerate(losses):
        print ii, li
        if p1dct[ii] != 0:
            assert li == p0dct[ii] ** 2
        else:
            assert li == 1

def test_vectorize_simple():
    N = as_apply(15)

    p0 = hp_uniform('p0', 0, 1)
    loss = p0 ** 2
    print loss
    expr_idxs = scope.range(N)
    vh = VectorizeHelper(loss, expr_idxs, build=True)
    vloss = vh.v_expr

    full_output = as_apply([vloss,
        vh.idxs_by_label(),
        vh.vals_by_label()])
    fo2 = replace_repeat_stochastic(full_output)

    new_vc = recursive_set_rng_kwarg(
            fo2,
            as_apply(np.random.RandomState(1)),
            )

    #print new_vc
    losses, idxs, vals = rec_eval(new_vc)
    print 'losses', losses
    print 'idxs p0', idxs['p0']
    print 'vals p0', vals['p0']
    p0dct = dict(zip(idxs['p0'], vals['p0']))
    for ii, li in enumerate(losses):
        assert p0dct[ii] ** 2 == li

def n_arms(N=2):
    """
    Each arm yields a reward from a different Gaussian.

    The correct arm is arm 0.

    """
    x = hp_choice('x', [0, 1])
    reward_mus = as_apply([-1] + [0] * (N - 1))
    reward_sigmas = as_apply([1] * N)
    return {'loss': scope.normal(reward_mus[x], reward_sigmas[x]),
            'loss_variance': 1.0}

    def __init__(self, bandit, seed=seed, cmd=None, workdir=None):
        self.bandit = bandit
        self.seed = seed
        self.rng = np.random.RandomState(self.seed)
        self.cmd = cmd
        self.workdir = workdir
        self.s_new_ids = pyll.Literal('new_ids')  # -- list at eval-time
        before = pyll.dfs(self.bandit.expr)
        # -- raises exception if expr contains cycles
        pyll.toposort(self.bandit.expr)
        vh = self.vh = VectorizeHelper(self.bandit.expr, self.s_new_ids)
        # -- raises exception if v_expr contains cycles
        pyll.toposort(vh.v_expr)

        idxs_by_label = vh.idxs_by_label()
        vals_by_label = vh.vals_by_label()
        after = pyll.dfs(self.bandit.expr)
        # -- try to detect if VectorizeHelper screwed up anything inplace
        assert before == after
        assert set(idxs_by_label.keys()) == set(vals_by_label.keys())
        assert set(idxs_by_label.keys()) == set(self.bandit.params.keys())

        # -- make the graph runnable and SON-encodable
        # N.B. operates inplace
        self.s_idxs_vals = recursive_set_rng_kwarg(
                scope.pos_args(idxs_by_label, vals_by_label),
                pyll.as_apply(self.rng))

        # -- raises an exception if no topological ordering exists
        pyll.toposort(self.s_idxs_vals)

 def test2(self):
     self.expr = as_apply(dict(p0=one_of(0, 1)))
     self.wanted = [
             [('p0.randint', [0], [0])], [('p0.randint', [1], [1])],
             [('p0.randint', [2], [0])], [('p0.randint', [3], [0])],
             [('p0.randint', [4], [0])]]
     self.foo()

 def evaluate(self, config, ctrl):
     memo = self.memo_from_config(config)
     memo[self.pyll_ctrl] = ctrl
     if self.init_pyll_memo:
         memo = self.init_pyll_memo(memo, config=config, ctrl=ctrl)
     if self.rng is not None and not self.installed_rng:
         # -- N.B. this modifies the expr graph in-place
         #    XXX this feels wrong
         self.expr = recursive_set_rng_kwarg(self.expr,
             pyll.as_apply(self.rng))
         self.installed_rng = True
     try:
         # -- the "work" of evaluating `config` can be written
         #    either into the pyll part (self.expr)
         #    or the normal Python part (self.fn)
         pyll_rval = pyll.rec_eval(self.expr, memo=memo)
         rval = self.fn(pyll_rval)
     except Exception, e:
         n_match = 0
         for match, match_pair in self.exceptions:
             if match(e):
                 rval = match_pair(e)
                 logger.info('Caught fn exception %s' % str(rval))
                 n_match += 1
                 break
         if n_match == 0:
             raise

def test_repeatable():
    u = scope.uniform(0, 1)
    aa = as_apply(dict(u=u, n=scope.normal(5, 0.1), l=[0, 1, scope.one_of(2, 3), u]))
    dd1 = sample(aa, np.random.RandomState(3))
    dd2 = sample(aa, np.random.RandomState(3))
    dd3 = sample(aa, np.random.RandomState(4))
    assert dd1 == dd2
    assert dd1 != dd3

def test_recursive_set_rng_kwarg():
    uniform = scope.uniform
    a = as_apply([uniform(0, 1), uniform(2, 3)])
    rng = np.random.RandomState(234)
    recursive_set_rng_kwarg(a, rng)
    print a
    val_a = rec_eval(a)
    assert 0 < val_a[0] < 1
    assert 2 < val_a[1] < 3

 def test5(self):
     p0 = uniform(0, 1)
     self.expr = as_apply(dict(p0=p0, p1=p0))
     self.wanted = [[('p0', [0], [0.69646918559786164])], [('p0', [1],
                 [0.28613933495037946])], [('p0', [2],
                     [0.22685145356420311])], [('p0', [3],
                         [0.55131476908289123])], [('p0', [4],
                             [0.71946896978556307])]]
     self.foo()

 def __init__(self):
     d = dict(
         # -- alphabetical order
         lu = scope.loguniform(-2, 2),
         qlu = scope.qloguniform(np.log(0.01), np.log(20), 2),
         qu = scope.quniform(-4.999, 5, 1),
         u = scope.uniform(0, 10),
         )
     base.Bandit.__init__(self, as_apply(d))

 def test1(self):
     self.expr = as_apply(dict(p0=normal(0, 1)))
     self.wanted = [
             [('p0', [0], [-1.0856306033005612])],
             [('p0', [1], [0.99734544658358582])],
             [('p0', [2], [0.28297849805199204])],
             [('p0', [3], [-1.506294713918092])],
             [('p0', [4], [-0.57860025196853637])]]
     self.foo()

 def test6(self):
     p0 = uniform(0, 1)
     self.expr = as_apply(dict(p0=p0, p1=normal(p0, 1)))
     self.wanted = [
         [('p0', [0], [0.69646918559786164]), ('p1', [0], [-0.25562802126346051])],
         [('p0', [1], [0.55131476908289123]), ('p1', [1], [-0.19412629039976703])],
         [('p0', [2], [0.71946896978556307]), ('p1', [2], [1.0415750381251847])],
         [('p0', [3], [0.68482973858486329]), ('p1', [3], [0.63331201764547818])],
         [('p0', [4], [0.48093190148436094]), ('p1', [4], [-1.1383681635523848])]]
     self.foo()

def test_sample():
    u = scope.uniform(0, 1)
    aa = as_apply(dict(u=u, n=scope.normal(5, 0.1), l=[0, 1, scope.one_of(2, 3), u]))
    print aa
    dd = sample(aa, np.random.RandomState(3))
    assert 0 < dd["u"] < 1
    assert 4 < dd["n"] < 6
    assert dd["u"] == dd["l"][3]
    assert dd["l"][:2] == (0, 1)
    assert dd["l"][2] in (2, 3)

def test_fg11_top_bandit():
    L = lfw.FG11Bandit()
    config = stochastic.sample(L.template, np.random.RandomState(0))
    config['decisions'] = None
    config['slm'] = stochastic.sample(pyll.as_apply(params.fg11_top),
            np.random.RandomState(0))
    config['comparison'] = 'sqrtabsdiff'
    rec = L.evaluate(config, hyperopt.base.Ctrl(None))
    assert np.abs(rec['loss'] - .194) < 1e-2
    return rec

def config0():
    p0 = scope.uniform(0, 1)
    p1 = scope.uniform(2, 3)
    p2 = scope.one_of(-1, p0)
    p3 = scope.one_of(-2, p1)
    p4 = 1
    p5 = [3, 4, p0]
    d = locals()
    del d['p1'] # -- don't sample p1 all the time
    s = as_apply(d)
    return s

def config0():
    p0 = scope.uniform(0, 1)
    p1 = scope.uniform(2, 3)
    p2 = scope.one_of(-1, p0)
    p3 = scope.one_of(-2, p1)
    p4 = 1
    p5 = [3, 4, p0]
    p6 = scope.one_of(-3, p1)
    d = locals()
    d['p1'] = None # -- don't sample p1 all the time, only if p3 says so
    s = as_apply(d)
    return s

 def build_vals(self):
     for node in self.dfs_nodes:
         if node.name == "literal":
             n_times = scope.len(self.idxs_memo[node])
             vnode = scope.asarray(scope.repeat(n_times, node))
         elif node in self.choice_memo:
             # -- choices are natively vectorized
             choices = self.choice_memo[node]
             self.vals_memo[choices] = choices
             # -- this stitches together the various sub-graphs
             #    to define the original node
             vnode = scope.vchoice_merge(self.idxs_memo[node], self.choice_memo[node])
             vnode.pos_args.extend(
                 [as_apply([self.idxs_memo[inode], self.vals_memo[inode]]) for inode in node.pos_args]
             )
         else:
             vnode = scope.idxs_map(self.idxs_memo[node], node.name)
             vnode.pos_args.extend(node.pos_args)
             vnode.named_args.extend(node.named_args)
             for arg in node.inputs():
                 vnode.replace_input(arg, as_apply([self.idxs_memo[arg], self.vals_memo[arg]]))
         self.vals_memo[node] = vnode

    def evaluate(self, config, ctrl, attach_attachments=True):
        memo = self.memo_from_config(config)
        self.use_obj_for_literal_in_memo(ctrl, base.Ctrl, memo)
        if self.rng is not None and not self.installed_rng:
            # -- N.B. this modifies the expr graph in-place
            #    XXX this feels wrong
            self.expr = recursive_set_rng_kwarg(self.expr,
                pyll.as_apply(self.rng))
            self.installed_rng = True
        if self.pass_expr_memo_ctrl:
            rval = self.fn(
                    expr=self.expr,
                    memo=memo,
                    ctrl=ctrl,
                    *self.args)
        else:
            # -- the "work" of evaluating `config` can be written
            #    either into the pyll part (self.expr)
            #    or the normal Python part (self.fn)
            pyll_rval = pyll.rec_eval(self.expr, memo=memo,
                    print_node_on_error=self.rec_eval_print_node_on_error)
            rval = self.fn(pyll_rval, *self.args)

        if isinstance(rval, (float, int, np.number)):
            dict_rval = {'loss': rval}
        elif isinstance(rval, (dict,)):
            dict_rval = rval
            if 'loss' not in dict_rval:
                raise ValueError('dictionary must have "loss" key',
                        dict_rval.keys())
        else:
            raise TypeError('invalid return type (neither number nor dict)', rval)

        if dict_rval['loss'] is not None:
            # -- fail if cannot be cast to float
            dict_rval['loss'] = float(dict_rval['loss'])

        dict_rval.setdefault('status', base.STATUS_OK)
        if dict_rval['status'] not in base.STATUS_STRINGS:
            raise ValueError('invalid status string', dict_rval['status'])

        if attach_attachments:
            attachments = dict_rval.pop('attachments', {})
            for key, val in attachments.items():
                ctrl.attachments[key] = val

        # -- don't do this here because SON-compatibility is only a requirement
        #    for trials destined for a mongodb. In-memory rvals can contain
        #    anything.
        #return base.SONify(dict_rval)
        return dict_rval