import itertools
import unittest


def all_friend_pairs(seating_arrangement):

    """ Determines all pairs of friends in a given seating arrangement.

    :param seating_arrangement: List of ints representing the seating arrangement
    around a circular table, with each int corresponding to the id of an individual.
    :return: set of pairs of ints corresponding to the id of individuals, such that
    for each pair (p_a, p_b), we have p_a < p_b, and p_a next to p_b in the list.
    """

    # we take each element and its neighbour, and sort the resulting tuple.
    fp = set(
        map(
            lambda x: (min(x), max(x)),
            zip(
                seating_arrangement,
                itertools.islice(itertools.cycle(seating_arrangement), 1, None),
            ),
        )
    )

    return fp


def get_pair(current_seating_arrangement, i, j):

    """ Returns a sorted tuple (p_a, p_b) of the labels corresponding to the people
    in seats i and j, where p_a < p_b

    :param current_seating_arrangement: list of integers representing the seating
    arrangement around a circular table
    :param i: index of one integer in list
    :param j: index of other integer in list
    :return: sorted tuple (p_a, p_b) of the elements at indices i and j in the list
    """
    n = len(current_seating_arrangement)

    p_i = current_seating_arrangement[i % n]
    p_j = current_seating_arrangement[j % n]

    if p_i < p_j:
        return (p_i, p_j)
    else:
        return (p_j, p_i)


def new_friend_pairs_v1(current_seating_arrangement, existing_friend_pairs, swap):

    """ Given a list of existing friends, determines the new friendships that would
    be formed when a pair of people in the current seating arrangement.

    :param current_seating_arrangement: list of integers representing the seating
    arrangement around a circular table
    :param existing_friend_pairs: list of sorted pairs of integers representing pairs of
    individuals who are already friends
    :param swap: pair of integers representing the positions (not IDs!) of the people in
    current_seating_arrangement we want to swap
    :return: list of pairs of integers representing pairs of individuals who
    would become friends by this new sitting arrangement
    """

    (i, j) = swap

    # carry out shallow copy of elements
    new_seating_arrangement = current_seating_arrangement[:]
    new_seating_arrangement[j] = current_seating_arrangement[i]
    new_seating_arrangement[i] = current_seating_arrangement[j]

    actual_new_friend_pairs = filter(
        lambda x: x not in existing_friend_pairs,
        all_friend_pairs(new_seating_arrangement),
    )

    return actual_new_friend_pairs


def new_friend_pairs(current_seating_arrangement, existing_friend_pairs, swap):

    (i, j) = swap
    candidate_new_friend_pairs = list()

    candidate_new_friend_pairs.append(get_pair(current_seating_arrangement, i, j - 1))
    candidate_new_friend_pairs.append(get_pair(current_seating_arrangement, i, j + 1))
    candidate_new_friend_pairs.append(get_pair(current_seating_arrangement, j, i - 1))
    candidate_new_friend_pairs.append(get_pair(current_seating_arrangement, j, i + 1))

    # we want to make sure that a candidate new_friend_pair is not an existing friend_pair.
    # the other thing is a hack to make sure we don't pretend a person is a friend with him/herself.
    actual_new_friend_pairs = list(
        filter(
            lambda x: x not in existing_friend_pairs and x[0] != x[1],
            candidate_new_friend_pairs,
        )
    )

    return actual_new_friend_pairs


def new_friend_pairs_v3(current_seating_arrangement, existing_friend_pairs, swap):

    (i, j) = swap
    candidate_new_friend_pairs = list()

    for (u, v) in [(i, j - 1), (i, j + 1), (j, i - 1), (j, i + 1)]:
        if u != v:
            pair = get_pair(current_seating_arrangement, u, v)
            if pair not in existing_friend_pairs:
                candidate_new_friend_pairs.append(pair)

    return candidate_new_friend_pairs


def num_new_friend_pairs(current_seating_arrangement, existing_friend_pairs, swap):

    """ Determines the number of new friend pairs generated by the current_swap.

    :param current_seating_arrangement: list of ints representing the seating
    arrangement around a circular table.
    :param existing_friend_pairs: set of pairs of integers representing pairs of
    individuals who are already friends
    :param swap: pair of integers representing the positions (not IDs!) of the people in
    current_seating_arrangement we want to swap
    :return: The number of new friend pairs generated by the swap.
    """

    return len(
        new_friend_pairs(current_seating_arrangement, existing_friend_pairs, swap)
    )


class TestNumNewFriendPairs(unittest.TestCase):

    # Tests num_new_friend_pairs

    def setUp(self):
        self.sa = range(10)
        self.fp = all_friend_pairs(self.sa)

    def test_self_swap(self):
        self.assertEqual(num_new_friend_pairs(self.sa, self.fp, (0, 0)), 0)

    def test_neighbor_swap(self):
        self.assertEqual(num_new_friend_pairs(self.sa, self.fp, (0, 1)), 2)

    def test_far_swap(self):
        self.assertEqual(num_new_friend_pairs(self.sa, self.fp, (0, 5)), 4)
