from functools import partial
import numpy as np
import unittest
from sklearn import datasets
from hyperopt import hp, Trials, fmin, tpe, rand
import hyperopt.pyll.stochastic


class TestPChoice(unittest.TestCase):
    def test_basic(self):

        space = hp.pchoice('naive_type',
                [(.14, 'gaussian'),
                 (.02, 'multinomial'),
                 (.84, 'bernoulli')])
        a, b, c = 0, 0, 0
        rng = np.random.RandomState(123)
        for i in range(0, 1000):
            nesto = hyperopt.pyll.stochastic.sample(space, rng=rng)
            if nesto == 'gaussian':
                a += 1
            elif nesto == 'multinomial':
                b += 1
            elif nesto == 'bernoulli':
                c += 1
        print(a, b, c)
        assert a + b + c == 1000
        assert 120 < a < 160
        assert 0 < b < 40
        assert 800 < c < 900

    def test_basic2(self):
        space = hp.choice('normal_choice', [
            hp.pchoice('fsd',
                [(.1, 'first'),
                 (.8, 'second'),
                 (.1, 2)]),
            hp.choice('something_else', [10, 20])
        ])
        a, b, c = 0, 0, 0
        rng=np.random.RandomState(123)
        for i in range(0, 1000):
            nesto = hyperopt.pyll.stochastic.sample(space, rng=rng)
            if nesto == 'first':
                a += 1
            elif nesto == 'second':
                b += 1
            elif nesto == 2:
                c += 1
            elif nesto in (10, 20):
                pass
            else:
                assert 0, nesto
        print(a, b, c)
        assert b > 2 * a
        assert b > 2 * c


class TestSimpleFMin(unittest.TestCase):
    # test that that a space with a pchoice in it is
    # (a) accepted by various algos and
    # (b) handled correctly.
    #
    def setUp(self):
        self.space = hp.pchoice('a', [
            (.1, 0), 
            (.2, 1),
            (.3, 2),
            (.4, 3)])
        self.trials = Trials()

    def objective(self, a):
        return [1, 1, 1, 0 ][a]

    def test_random(self):
        # test that that a space with a pchoice in it is
        # (a) accepted by tpe.suggest and
        # (b) handled correctly.
        N = 50
        fmin(self.objective,
            space=self.space,
            trials=self.trials,
            algo=rand.suggest,
            max_evals=N)

        a_vals = [t['misc']['vals']['a'][0] for t in self.trials.trials]
        counts = np.bincount(a_vals)
        print counts
        assert counts[3] > N * .35
        assert counts[3] < N * .60

    def test_tpe(self):
        N = 50
        fmin(self.objective,
            space=self.space,
            trials=self.trials,
            algo=partial(tpe.suggest, n_startup_jobs=10),
            max_evals=N)

        a_vals = [t['misc']['vals']['a'][0] for t in self.trials.trials]
        counts = np.bincount(a_vals)
        print counts
        assert counts[3] > N * .7


@@ -1653,16 +1653,864 @@
  2 * c%0A%0A
+    def test_basic3(self):%0A        space = hp.pchoice('something', %5B%0A            (.2, hp.pchoice('number', %5B(.8, 2), (.2, 1)%5D)),%0A            (.8, hp.pchoice('number1', %5B(.7, 5), (.3, 6)%5D))%0A        %5D)%0A        a, b, c, d = 0, 0, 0, 0%0A        rng=np.random.RandomState(123)%0A        for i in range(0, 2000):%0A            nesto = hyperopt.pyll.stochastic.sample(space, rng=rng)%0A            if nesto == 2:%0A                a += 1%0A            elif nesto == 1:%0A                b += 1%0A            elif nesto == 5:%0A                c += 1%0A            elif nesto == 6:%0A                d += 1%0A            else:%0A                assert 0, nesto%0A        print(a, b, c, d)%0A        assert a + b + c + d == 2000%0A        assert 300 %3C a + b %3C 500%0A        assert 1500 %3C c + d %3C 1700%0A        assert a * .3 %3E b  # a * 1.2 %3E 4 * b%0A        assert c * 3 * 1.2 %3E d * 7%0A        %0A
 %0Aclass T

test case in unit testing