from collections import defaultdict, deque
from typing import List


class Vertex:
    def __init__(self, identifier):
        self.id = identifier
        self.neighbors = dict()  # id --> Vertex

    def neighbor_ids(self):
        return list(self.neighbors.keys())


class GraphList:
    def __init__(self, graph):
        self.vertices = dict()  # id --> Vertex
        # add the nodes themselves
        for node_id in range(len(graph)):
            self.vertices[node_id] = Vertex(node_id)
        # connect each node to its neighbors
        for node_id in range(len(graph)):
            for other_id in range(len(graph)):
                if graph[node_id][other_id] == 1:
                    other_node = self.vertices[other_id]
                    self.vertices[node_id].neighbors[other_id] = other_node


class Solution:
    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:
        """
        sq. adj. matrix --> graph ADT

        when a node is infected, it also infects all of its direct neighbors (if not already)

        spread stops when no more nodes can be infected

        inputs
        matrix = the network
        infected = ids of some nodes already infected

        Intuition:
            Goal - min the total # of nodes infected by the spread
                 - return the id of that node we could remove to min the spread

        Approach
            1. Play Out Multiple Spread:
                1. for each infected
                    a) H: compute the final # of infected nodes, excluding that node [TODO]
                    b) map the final num --> ids that lead to that number
                2. return the id of the init ---> that min the final #

                ---------------
                Helper - initial
                a) computing the final number of infected nodes
                    input -
                        list of ids,
                        adj matrix that shows connections between nodes

                    A: init a dict, assume all the nodes (repr by id) are NOT infected (False)
                    B: BFS through the infected nodes, count up all the infected
                         b) BFS through the infected nodes
                            c) mark them as infected in the dict
                            d) mark + enqueue their neighbors, if not already marked True
                    C: use the dict to find the final # of infected - lowest value in lowest key

            2. 1 DFS
                New Intutition:
                    1. as long as there is at least 1 infected node in a set of connected components (CC),
                        then every node in that set will be infected
                    2. and we can count the size of that set of that set towards
                        the total num of infected
                    3. if there is more than 1 initally infected (II) node in a set of CC,
                        then don't return any of those nodes as the answer
                        (it won't decrease the total # of nodes eventually infected)
                    4. init the answer as the min of the II nodes
                    5. Then try to order the nodes by:
                        - if they are the only node in their set of CC
                        - AND the size of the set of nodes they are connected with

                Approach:
                    A: init the answer as the min of the II nodes
                    B:
                        1. create graph adj. lists for each set of CC
                            - Graph obj --> knows the size of the CC
                                --> # of II
                            - Vertex objs ---> they know their id (helps w/ ans)
                        2. Do a DFS on the CC that each II is a part of
                            - w/ just 1 Vertex object, we still can get:
                                - size of the CC
                                - id of this node
                                - can tell if any other II are in the same set
                    C: Compute the lowest Vertex id amongst II in each CC
                    D: Sort the Vertex objs by their size, and
                        then sort by the lowest id


                BUT - you must consider what happens when there's more than 1 CC
                        ---> ALSO have to consider if that node is the only II node
                             in its CC



                [     0 1 2
                   0 [1,1,0], ------>        Vertex(0) --- Vertex(1)
                   1 [1,1,0],
                   2 [0,0,1]                        Vertex(2)
                ],
                initial = [0,1]
                -------->[
                    {
                        id: 0,
                        neighbors: { 1: True }
                    },
                    {
                        id: 1,
                        neighbors: { 0: True }
                    }

                    (size, id)
                    [(1, 1), (1, 0)]
                ]
                Output: 0


        Edge Cases:
            - multiple nodes that could min ---> return the smallest id

        ---------------------
        Test Run:
        graph = [
              0 1. 2
          0  [1,1, 0],
          1  [0,1, 0],
          2  [0,0, 1],
        ]

        initial = [0,1]

        exclude_id = 0

        final_num_given_excluded_id = {
            2: [0, 1]
        }

        -------
        Helper, 1st call:

        helper_1_initial = q = []
        ---> 1, 0

        infected = {
            0: False, ---> True
            1: False, --> True
            2: False
        }


        """
        """
        # A: compute how many infections each initial node cont.
        infections_per_initial = dict(zip(
                [node for node in initial], 
                [0 for _ in range(len(initial))]
        ))
        # B: BFS + prevent cycles
        infected_node_ids = dict(zip(
                list(range(len(graph))),
                [False for _ in range(len(graph))]
        ))
        q = deque(initial)
        # C: start BFS
        while len(q) > 0:
            infected_id = q.popleft()
            # c) mark them as infected in the dict
            infected_node_ids[infected_id] = True
            # d) mark + enqueue their neighbors, if not already marked True 
            for other_id in range(len(graph)):
                if graph[other_id][infected_id] == 1:
                    if infected_node_ids[other_id] is False:
                        q.append(other_id)
                        # also update the cont of the initial
                        if infected_id in infections_per_initial:
                            infections_per_initial[infected_id] += 1
        # D: return the smallest id amongast those that caused highest infections
        highest_contrib = max(infections_per_initial.values())
        lowest_initial_id = min([
            node_id for node_id in infections_per_initial.keys()
            if infections_per_initial[node_id] == highest_contrib
        ])
        return lowest_initial_id
        
        """

        def compute_final_infected(excluded_id):
            # A: init a dict, assume all the nodes (repr by id) are NOT infected (False)
            infected_node_ids = dict(
                zip(list(range(len(graph))), [False for _ in range(len(graph))])
            )
            # B: BFS through the infected nodes, count up all the infected
            q = deque([val for val in initial if val != excluded_id])
            # b) BFS through the infected nodes
            while len(q) > 0:
                infected_id = q.popleft()
                # c) mark them as infected in the dict
                infected_node_ids[infected_id] = True
                # d) mark + enqueue their neighbors, if not already marked True
                for other_id in range(len(graph)):
                    if graph[other_id][infected_id] == 1:
                        if infected_node_ids[other_id] is False:
                            q.append(other_id)
            # C: use the dict to find the final # of infected - lowest value in lowest key
            count = sum(list(infected_node_ids.values()))
            return count

        # init dict
        final_num_given_excluded_id = defaultdict(list)
        # 1. for each infected
        for excluded_id in initial:
            # a) H: compute the final # of infected nodes, excluding that node [TODO]
            final_infected = compute_final_infected(excluded_id)
            # b) map the final num --> ids that lead to that number
            final_num_given_excluded_id[final_infected].append(excluded_id)
        # 2. return the id of the init ---> that min the final #
        minimum_infected = min(list(final_num_given_excluded_id.keys()))
        return min(final_num_given_excluded_id[minimum_infected])


"""
graph = [
      0 1.2 
   0 [1,1,0],
   1 [1,1,0],
   2 [0,0,1]
], 

initial = [0,1]
           ^

final_num_given_excluded_id = {
    2: [0, ]
}


excluded_id     remaining_initial = q      final_infected
0                   []                       2


infected_node_ids = {
    0: F --> T --> 1
    1: F --> T --> 1
    2: F --------> 0
}

infected_id = 1, 0     other_id = 0, 1, 2, 0, 1, 2
                        
"""
