from collections import Counter
import itertools
import random
import unittest
from example import reservoir_sample, vertex_cover_2approx, membership_benchmark


class Choices:
    def __init__(self, values):
        self.values = iter(values)
    def randrange(self, bound):
        value = next(self.values)
        assert 0 <= value < bound
        return value


class EngineeringTests(unittest.TestCase):
    def test_sample_edges_and_exact_small_distribution(self):
        self.assertEqual(reservoir_sample([], 3, random.Random(1)), [])
        self.assertEqual(reservoir_sample(range(3), 9, random.Random(1)), [0, 1, 2])
        self.assertEqual(reservoir_sample(range(3), 0, random.Random(1)), [])
        counts = Counter()
        for choices in itertools.product(range(2), range(3)):
            counts.update(reservoir_sample(range(3), 1, Choices(choices)))
        self.assertEqual(counts, {0: 2, 1: 2, 2: 2})
        with self.assertRaises(ValueError):
            reservoir_sample([], -1, random.Random(1))

    def test_cover_and_approximation_bound(self):
        possible = list(itertools.combinations(range(4), 2))
        for mask in range(1 << len(possible)):
            edges = [edge for i, edge in enumerate(possible) if mask & (1 << i)]
            cover = vertex_cover_2approx(edges)
            self.assertTrue(all(a in cover or b in cover for a, b in edges))
            optimum = min((len(picked) for size in range(5)
                           for picked in itertools.combinations(range(4), size)
                           if all(a in picked or b in picked for a, b in edges)), default=0)
            self.assertLessEqual(len(cover), 2 * optimum)

    def test_benchmark_checks_meaning_not_speed(self):
        result = membership_benchmark(16, repeats=2)
        self.assertEqual(result['list']['hits'], result['set']['hits'])
        self.assertGreaterEqual(result['list']['median_seconds'], 0)


if __name__ == '__main__':
    unittest.main()
