#!/usr/bin/python3

import unittest
import numpy as np
import fairdiv
import generation
import solving

INSTANCE1 = fairdiv.Instance(np.array([[9, 8, 2, 1], [2, 5, 1, 4]]))
INSTANCE2 = fairdiv.Instance(np.array([[8, 2, 1], [5, 1, 5]]))
INSTANCE3 = fairdiv.Instance(np.array([[2, 3, 3, 2], [2, 3, 4, 1], [0, 4, 2, 4]]))

ALLOCATION1 = fairdiv.Allocation(
    INSTANCE1,
    np.array([[True, False, False, True], [False, True, True, False]]))

ALLOCATION2_1 = fairdiv.Allocation(
    INSTANCE2,
    np.array([[True, False, False], [False, True, True]]))
ALLOCATION2_2 = fairdiv.Allocation(
    INSTANCE2,
    np.array([[False, True, False], [True, False, True]]))
ALLOCATION2_3 = fairdiv.Allocation(
    INSTANCE2,
    np.array([[True, False, True], [False, True, False]]))

ALLOCATION3_1 = fairdiv.Allocation(
    INSTANCE3,
    np.array([[True, False, False, True], [False, False, True, False], [False, True, False, False]]))
ALLOCATION3_2 = fairdiv.Allocation(
    INSTANCE3,
    np.array([[False, False, False, True], [True, False, True, False], [False, True, False, False]]))



class TestInstances(unittest.TestCase):
    def setUp(self):
        self.instance1 = INSTANCE1
        self.instance3 = INSTANCE3

    def test_w_max(self):
        actual = self.instance1.w_max(0)
        expected = 9
        self.assertEqual(actual, expected)
        actual = self.instance1.w_max(1)
        expected = 5
        self.assertEqual(actual, expected)

    def test_best(self):
        actual = self.instance1.best(0)
        expected = [0]
        self.assertEqual(actual, expected)
        actual = self.instance1.best(1)
        expected = [1]
        self.assertEqual(actual, expected)

    def test_restricted_w_max(self):
        actual = self.instance1.w_max(0, [False, False, True, True])
        expected = 2
        self.assertEqual(actual, expected)
        actual = self.instance1.w_max(1, [True, False, True, True])
        expected = 4
        self.assertEqual(actual, expected)

    def test_restricted_best(self):
        actual = self.instance1.best(0, [False, False, True, True])
        expected = [2]
        self.assertEqual(actual, expected)
        actual = self.instance1.best(1, [True, False, True, True])
        expected = [3]
        self.assertEqual(actual, expected)

    def test_all_allocations(self):
        actual = sum(1 for _ in self.instance1.allocations())
        expected = 16
        self.assertEqual(actual, expected)

    def test_utility(self):
        expected = 9
        actual = self.instance1.utility(0, (True, False, False, False))
        self.assertEqual(actual, expected)
        expected = 6
        actual = self.instance1.utility(1, (False, True, True, False))
        self.assertEqual(actual, expected)

    def test_better_than(self):
        expected = [
            (False, True, True, True),
            (True, False, True, False),
            (True, False, True, True),
            (True, True, False, False),
            (True, True, False, True),
            (True, True, True, False),
            (True, True, True, True)
        ]
        count = 0
        for share in solving.better_than(self.instance1, 0, 10):
            count += 1
            self.assertIn(share, expected)
        self.assertEqual(count, len(expected))

    def test_max_min_share(self):
        actual = solving.max_min_share(self.instance3, 0)
        expected = 3
        self.assertEqual(actual, expected)
        actual = solving.max_min_share(self.instance3, 1)
        expected = 3
        self.assertEqual(actual, expected)
        actual = solving.max_min_share(self.instance3, 2)
        expected = 2
        self.assertEqual(actual, expected)
        actual = solving.max_min_share(self.instance3, 0, 2)
        expected = 5
        self.assertEqual(actual, expected)
        actual = solving.max_min_share(self.instance3, 1, 2)
        expected = 5
        self.assertEqual(actual, expected)
        actual = solving.max_min_share(self.instance3, 2, 2)
        expected = 4
        self.assertEqual(actual, expected)

    def test_min_max_share(self):
        actual = solving.min_max_share(self.instance3, 0)
        expected = 4
        self.assertEqual(actual, expected)
        actual = solving.min_max_share(self.instance3, 1)
        expected = 4
        self.assertEqual(actual, expected)
        actual = solving.min_max_share(self.instance3, 2)
        expected = 4
        self.assertEqual(actual, expected)
        actual = solving.min_max_share(self.instance3, 0, 2)
        expected = 5
        self.assertEqual(actual, expected)
        actual = solving.min_max_share(self.instance3, 1, 2)
        expected = 5
        self.assertEqual(actual, expected)
        actual = solving.min_max_share(self.instance3, 2, 2)
        expected = 6
        self.assertEqual(actual, expected)


class TestAllocations(unittest.TestCase):
    def setUp(self):
        self.allocation1 = ALLOCATION1
        self.allocation2_1 = ALLOCATION2_1
        self.allocation2_2 = ALLOCATION2_2
        self.allocation2_3 = ALLOCATION2_3
        self.allocation3_1 = ALLOCATION3_1
        self.allocation3_2 = ALLOCATION3_2

    def test_proportional(self):
        actual = self.allocation1.is_proportional()
        expected = True
        self.assertEqual(actual, expected)

    def test_envy_free(self):
        actual = self.allocation1.is_envy_free()
        expected = True
        self.assertEqual(actual, expected)

    def test_swap_optimal(self):
        actual = self.allocation1.is_swap_optimal()
        expected = False
        self.assertEqual(actual, expected)

    def test_sequenceable(self):
        actual = self.allocation1.is_sequenceable()
        expected = False
        self.assertEqual(actual, expected)

    def test_pareto_optimal(self):
        actual = self.allocation1.is_pareto_optimal()
        expected = False
        self.assertEqual(actual, expected)

    def test_sequences(self):
        actual = self.allocation2_1.sequence()
        expected = [[1, 0, 1], [0, 1, 1]]
        self.assertIn(actual, expected)
        actual = self.allocation2_2.sequence()
        expected = [1, 0, 1]
        self.assertEqual(actual, expected)
        actual = self.allocation2_3.sequence()
        expected = None
        self.assertEqual(actual, expected)

    def test_ceei(self):
        actual = self.allocation3_1.is_ceei()
        self.assertTrue(actual)
        actual = self.allocation3_2.is_ceei()
        self.assertFalse(actual)

    def test_max_min_share(self):
        actual = self.allocation3_1.is_max_min_share()
        self.assertTrue(actual)
        actual = self.allocation3_2.is_max_min_share()
        self.assertFalse(actual)

    def tearDown(self):
        pass


class TestGeneration(unittest.TestCase):
    def setUp(self):
        self.nb_tests = 10
        self.nb_agents = 3
        self.nb_objects = 10
        self.max_weight = 100

    def test_normalization(self):
        for _ in range(self.nb_tests):
            matrix = generation.generate_uniform(self.nb_agents,
                                                 self.nb_objects,
                                                 self.max_weight)
            for agent in range(self.nb_agents):
                for obj in range(self.nb_objects):
                    self.assertLessEqual(matrix[agent][obj], self.max_weight)
                    self.assertLessEqual(0, matrix[agent][obj])
            sums = matrix.sum(axis=1)
            sums -= sums.max()
            self.assertTrue(np.array_equal(
                sums, np.zeros(self.nb_agents, dtype=int)))

    def tearDown(self):
        pass


if __name__ == "__main__":
    import nose
    import sys

    MODULE_NAME = sys.modules[__name__].__file__

    nose.run(argv=[sys.argv[0], MODULE_NAME, '-v'] + sys.argv[1:])
