from collections import deque
import random

class User:
    def __init__(self, name):
        self.name = name

class SocialGraph:
    def __init__(self):
        self.lastID = 0
        self.users = {}
        self.friendships = {}

    def addFriendship(self, userID, friendID):
        """
        Creates a bi-directional friendship
        """
        if userID == friendID:
            print("WARNING: You cannot be friends with yourself")
        elif friendID in self.friendships[userID] or userID in self.friendships[friendID]:
            print("WARNING: Friendship already exists")
        else:
            self.friendships[userID].add(friendID)
            self.friendships[friendID].add(userID)

    def addUser(self, name):
        """
        Create a new user with a sequential integer ID
        """
        self.lastID += 1  # automatically increment the ID to assign the new user
        user = User(name)
        self.users[self.lastID] = user
        self.friendships[self.lastID] = set()
        # return user

    def populateGraph(self, numUsers, avgFriendships):
        """
        Takes a number of users and an average number of friendships
        as arguments

        Creates that number of users and a randomly distributed friendships
        between those users.

        The number of users must be greater than the average number of friendships.
        """
        # Reset graph
        self.lastID = 0
        self.users = {}
        self.friendships = {}

        # Add users
        while(self.lastID < numUsers):
            user = self.addUser(f"Person{self.lastID + 1}")

        # list containing user ids.
        rolodex = list(self.users.keys())

        # Add Friendships
        for userID, user in self.users.items():

            # creates copy of userID and shuffles it
            shuffled = rolodex.copy()
            random.shuffle(shuffled)

            # generates random number of friends for user
            num_friends = random.randint(0, avgFriendships * 2)

            for count in range(num_friends):
                friendID = shuffled.pop()
                if friendID != userID and userID < friendID:
                    self.addFriendship(userID, friendID)


    def getAllSocialPaths(self, userID):
        """
        Takes a user's userID as an argument

        Returns a dictionary containing every user in that user's
        extended network with the shortest friendship path between them.

        The key is the friend's ID and the value is the path.
        """
        visited = {userID:[userID]}

        queue = deque([userID])

        while(len(queue) > 0):
            personID = queue.popleft()
            friendIDs = self.friendships[personID]
            
            for friendID in friendIDs:
                # add to queue 
                
                if friendID not in visited:
                    queue.append(friendID)  
                    if personID in visited:

                        path = []
                        path = visited[personID].copy()
                        path.append(friendID)
                
                        visited[friendID] = path

        return visited


if __name__ == '__main__':
    sg = SocialGraph()
    sg.populateGraph(10, 2)
    print(sg.friendships)
    connections = sg.getAllSocialPaths(1)
    print(connections)

# sg = SocialGraph()
# sg.populateGraph(10,2)

# print(sg.friendships)
# connections = sg.getAllSocialPaths(1)
# print(connections)

