#!/usr/bin/env python3
import os
import lab
import sys
import pickle
import random
import unittest


TEST_DIRECTORY = os.path.dirname(__file__)

class Test00_Tiny(unittest.TestCase):
    def setUp(self):
        """ Load actor/movie database """
        filename = 'resources/tiny.pickle'
        with open(filename, 'rb') as f:
            self.data = pickle.load(f)

    def test_acted_together(self):
        actor1 = 4724
        actor2 = 1532
        result = lab.acted_together(self.data, actor1, actor2)
        self.assertTrue(result)

        actor1 = 4724
        actor2 = 4025
        result = lab.acted_together(self.data, actor1, actor2)
        self.assertFalse(result)

    def test_bacon_number_0(self):
        result = lab.actors_with_bacon_number(self.data, 0)
        expected = {4724}
        self.assertEqual(expected, result)

    def test_bacon_number_1(self):
        result = lab.actors_with_bacon_number(self.data, 1)
        expected = {1532, 2876}
        self.assertEqual(expected, result)

    def test_bacon_number_2(self):
        result = lab.actors_with_bacon_number(self.data, 2)
        expected = {1640}
        self.assertEqual(expected, result)

    def test_bacon_number_3(self):
        result = lab.actors_with_bacon_number(self.data, 3)
        expected = set()
        self.assertEqual(expected, result)

    def test_bacon_path(self):
        result = lab.bacon_path(self.data, 1640)
        expected = [4724, 2876, 1640]
        self.assertEqual(expected, result)


class Test01_ActedTogether(unittest.TestCase):
    def setUp(self):
        """ Load actor/movie database """
        filename = 'resources/small.pickle'
        with open(filename, 'rb') as f:
            self.data = pickle.load(f)

    def test_01(self):
        # Simple test, two actors who acted together
        actor1 = 4724
        actor2 = 9210
        self.assertTrue(lab.acted_together(self.data, actor1, actor2))

    def test_02(self):
        # Simple test, two actors who had not acted together
        actor1 = 4724
        actor2 = 16935
        self.assertFalse(lab.acted_together(self.data, actor1, actor2))

    def test_03(self):
        # Simple test, same actor
        actor1 = 4724
        actor2 = 4724
        self.assertTrue(lab.acted_together(self.data, actor1, actor2))


class Test02_BaconNumber(unittest.TestCase):
    def setUp(self):
        """ Load actor/movie database """
        filename = 'resources/small.pickle'
        with open(filename, 'rb') as f:
            self.data = pickle.load(f)

    def test_04(self):
        # Actors with Bacon number of 2
        n = 2
        expected = {1640, 1811, 2115, 2283, 2561, 2878, 3085, 4025, 4252, 4765,
                    6541, 9827, 11317, 14104, 16927, 16935, 19225, 33668, 66785,
                    90659, 183201, 550521, 1059002, 1059003, 1059004, 1059005,
                    1059006, 1059007, 1232763}

        first_result = lab.actors_with_bacon_number(self.data, n)
        self.assertTrue(isinstance(first_result, set))
        self.assertEqual(first_result, expected)

        second_result = lab.actors_with_bacon_number(self.data, n)
        self.assertTrue(isinstance(second_result, set))
        self.assertEqual(second_result, expected)

    def test_05(self):
        # Actors with Bacon number of 3
        n = 3
        expected = {52, 1004, 1248, 2231, 2884, 4887, 8979, 10500, 12521,
                    14792, 14886, 15412, 16937, 17488, 19119, 19207, 19363,
                    20853, 25972, 27440, 37252, 37612, 38351, 44712, 46866,
                    46867, 48576, 60062, 75429, 83390, 85096, 93138, 94976,
                    109625, 113777, 122599, 126471, 136921, 141458, 141459,
                    141460, 141461, 141495, 146634, 168638, 314092, 349956,
                    558335, 572598, 572599, 572600, 572601, 572602, 572603,
                    583590, 931399, 933600, 1086299, 1086300, 1168416, 1184797,
                    1190297, 1190298, 1190299, 1190300}
        first_result = lab.actors_with_bacon_number(self.data, n)
        self.assertTrue(isinstance(first_result, set))
        self.assertEqual(first_result, expected)

        second_result = lab.actors_with_bacon_number(self.data, n)
        self.assertTrue(isinstance(second_result, set))
        self.assertEqual(second_result, expected)

    def test_06(self):
        # random large Bacon number
        N = random.randint(50, 100)
        k = random.randint(7, 30)
        self.assertEqual(len(lab.actors_with_bacon_number(make_bacon_tree(N, k), N)), k)

    def test_07(self):
        # random graph, Bacon number with no people
        N = random.randint(5, 10)
        k = random.randint(4, 7)
        self.assertEqual(len(lab.actors_with_bacon_number(make_bacon_tree(N, k), int(1e20))), 0)
        self.assertEqual(len(lab.actors_with_bacon_number(make_bacon_tree(N, k), int(1e20))), 0)


class Test03_BaconPath(unittest.TestCase):
    def setUp(self):
        """ Load actor/movie database """
        with open('resources/small.pickle', 'rb') as f:
            self.db_small = pickle.load(f)
        with open('resources/large.pickle', 'rb') as f:
            self.db_large = pickle.load(f)

    def test_08(self):
        # Bacon path, small database, path does not exist
        actor_id = 2876669
        expected = None

        first_result = lab.bacon_path(self.db_small, actor_id)
        self.assertEqual(first_result, expected)

        second_result = lab.bacon_path(self.db_small, actor_id)
        self.assertEqual(second_result, expected)

    def test_09(self):
        # Bacon path, small database, length of 3 (4 actors, 3 movies)
        actor_id = 46866
        len_expected = 3

        first_result = lab.bacon_path(self.db_small, actor_id)
        len_first_result = -1 if first_result is None else len(first_result)-1
        second_result = lab.bacon_path(self.db_small, actor_id)
        len_second_result = -1 if second_result is None else len(second_result)-1

        self.assertTrue(valid_path(self.db_small, first_result))
        self.assertEqual(len_first_result, len_expected)
        self.assertEqual(first_result[0], 4724)
        self.assertEqual(first_result[-1], actor_id)

        self.assertTrue(valid_path(self.db_small, second_result))
        self.assertEqual(len_second_result, len_expected)
        self.assertEqual(second_result[0], 4724)
        self.assertEqual(second_result[-1], actor_id)

    def test_10(self):
        # Bacon path, large database, length of 2 (3 actors, 2 movies)
        actor_id = 1204
        len_expected = 2
        result = lab.bacon_path(self.db_large, actor_id)
        len_result = -1 if result is None else len(result)-1
        self.assertTrue(valid_path(self.db_large, result))
        self.assertEqual(len_result, len_expected)
        self.assertEqual(result[0], 4724)
        self.assertEqual(result[-1], actor_id)

    def test_11(self):
        # Bacon path, large database, length of 4 (5 actors, 4 movies)
        actor_id = 197897
        len_expected = 4
        result = lab.bacon_path(self.db_large, actor_id)
        len_result = -1 if result is None else len(result)-1
        self.assertTrue(valid_path(self.db_large, result))
        self.assertEqual(len_result, len_expected)
        self.assertEqual(result[0], 4724)
        self.assertEqual(result[-1], actor_id)

    def test_12(self):
        # Bacon path, large database, length of 6 (7 actors, 6 movies)
        actor_id = 1345462
        len_expected = 6
        result = lab.bacon_path(self.db_large, actor_id)
        # here, we compute the result twice, to test for mutation of the db
        result = lab.bacon_path(self.db_large, actor_id)
        len_result = -1 if result is None else len(result)-1
        self.assertTrue(valid_path(self.db_large, result))
        self.assertEqual(len_result, len_expected)
        self.assertEqual(result[0], 4724)
        self.assertEqual(result[-1], actor_id)

    def test_13(self):
        # Bacon path, large database, does not exist
        actor_id = 1204555
        expected = None
        result = lab.bacon_path(self.db_large, actor_id)
        self.assertEqual(result, expected)

def random_number_list(L, i=1):
    o = list(range(i*100000, i*100000+L))
    random.shuffle(o)
    return o


def valid_path(d, p):
    x = {frozenset(i) for i in d}
    return all(frozenset(i) in x for i in zip(p, p[1:]))


def make_bacon_tree(L, n=10):
    id_set = 2
    path = [4724] + random_number_list(L, i=1)
    n -= 1
    out = set((i,j) for i,j in zip(path, path[1:]))
    while n > 0:
        point = random.choice(range(len(path)-1))
        d = L - point
        if d == 0:
            continue
        newpath = random_number_list(d, i=id_set)
        p = [path[point]] + newpath
        out |= set((i,j) for i,j in zip(p, p[1:]))
        id_set += 1
        n -= 1
    return [(i, j) for i,j in out]


if __name__ == '__main__':
    res = unittest.main(verbosity=3, exit=False)
