# Import system services
import os
import sys

# The main package behind the graph data structures.
import networkx as nx
import random


class forest:
    """
    This class is created by an Election to track delegations that have happened. It stores delegations in a forest
    structure.
    """

    # The constructors for this class
    def __init__(self, n):
        # create empty graph with n nodes, add edges between them as we add delegations
        self.G = nx.DiGraph()
        self.G.add_nodes_from(list(range(n)))
        self.guru_tree = None   # arbitrary constructed tree in which all nodes are gurus

    def add_delegation(self, i, j):
        """
        Add a delegation from voter i to voter j.
        #TODO: This method could check for cycles and not add an edge that is not allowed
        :param i:
        :param j:
        :return: True or False depending on whether the delegation was made successfully
        """
        if self.get_guru_of_voter(j) == i:
            # i cannot delegate to someone who has i as a guru, would cause a cycle
            return False

        # remove any other delegations i is making
        i_edges = list(self.G.out_edges(i))
        self.G.remove_edges_from(i_edges)

        if i == j:
            # voter is becoming a guru, stop after removing all existing delegations
            return True

        # add new edge from i to j
        self.G.add_edge(i, j)

        # check if cycle exists
        try:
            cycle = nx.find_cycle(self.G, source=i)
        except nx.NetworkXNoCycle:
            cycle = False

        if cycle:
            # remove newly added edge
            # add back old delegations
            self.G.remove_edge(i, j)
            self.G.add_edges_from(i_edges)
            return False

        return True

    def make_random_delegation(self, delegator, voter_network, allow_guru_changes=False):
        """
        Have the given node make a random delegation. Based on the parameter, this may be a delegation to itself which
        causes the node to become (or remain) a guru. No guarantee that the result_accuracy will be different than the
        current delegation.
        :param delegator: index of the voter changing their delegation
        :param voter_network: The network the voters exist upon
        :param allow_guru_changes:
        :return:
        """
        neighbors = list(voter_network.neighbors(delegator))
        if allow_guru_changes:
            neighbors.append(delegator)
        # if allow_guru_changes and random.random() < 0.3:
        #     neighbors = [delegator]

        new_delegate = random.choice(neighbors)
        return self.add_delegation(delegator, new_delegate), new_delegate

    def get_gurus(self):
        """
        Return a list of the node indices that are gurus (have no outgoing edges). The result_accuracy has no ordering guarantee.
        :return:
        """
        sink_nodes = [node for node in self.G.nodes() if self.G.out_degree(node) == 0]
        return sink_nodes

    def get_delegators(self):
        """
        Return a list of the node indices that delegate (have an outgoing edge). The result_accuracy has no ordering guarantee.
        :return:
        """
        delegating_nodes = [node for node in self.G.nodes() if self.G.out_degree(node) == 1]
        return delegating_nodes

    def get_guru_of_voter(self, voter):
        """
        Return the index of the voter that is this voter's guru. Possibly the voter themself.
        :param voter:
        :return:
        """
        sinks = self.get_gurus()

        for sink in sinks:
            paths = []
            for vote_path in nx.all_simple_paths(self.G, voter, sink):
                paths.append(vote_path)
            if len(paths) > 0:
                break

        return paths

    def get_guru_weights(self):
        """
        Find the number of voters delegating to each guru, including itself.
        Return a dictionary mapping each guru index to their weight.
        # TODO: There may be a much faster way to do this
        :return:
        """
        # weights = []
        weights = dict()

        U = self.G.to_undirected()
        for node in self.get_gurus():
            w = nx.node_connected_component(U, node)
            weights[node] = len(w)
            # weights.append(len(w))

        return weights

    def get_guru_accuracies(self, all_accuracies):
        """
        Return a dictionary mapping the index of each guru to their accuracy.
        :param all_accuracies: List of accuracy of each voter
        :return:
        """
        return {i: all_accuracies[i] for i in self.get_gurus()}

    def get_delegation_list(self):
        """
        Return a list of length n where element i is the voter that voter i delegates to, or i if i is a guru.
        :return:
        """
        # TODO: Consider making this a parameter that is updated whenever a delegation changes
        delegations = []
        for i in range(len(self.G.nodes)):
            if self.G.out_degree(i) > 0:
                delegations.append(list(self.G.successors(i))[0])
            else:
                delegations.append(i)
        return delegations

    def create_guru_tree(self):
        """
        Create a binary tree containing all the current gurus as nodes.
        Pay no attention to the particular location of each nodes.
        :return:
        """

        gurus = self.get_gurus()
        self.guru_tree = nx.generators.classic.full_rary_tree(r=2, n=len(gurus), create_using=nx.DiGraph)

        # relabel nodes so indices match the guru indices
        mapping = {i: gurus[i] for i in range(len(gurus))}
        nx.relabel_nodes(self.guru_tree, mapping, copy=False)

    def accuracy_of_forest(self, all_accuracies):
        """
        Given accuracy of all voters in the entire network, determine the accuracy of system with the current
        delegations.
        :param all_accuracies:
        :return:
        """
        self.create_guru_tree()

        root = None
        for node in self.guru_tree:
            if len(list(self.guru_tree.predecessors(node))) == 0:
                root = node
                break
        guru_accuracies = self.get_guru_accuracies(all_accuracies)

        guru_weights = self.get_guru_weights()

        accs = self.subtree_accuracy(left_accs=None,
                                    right_accs=None,
                                    curr_weight=guru_weights[root],
                                    curr_idx=root,
                                    guru_weights=guru_weights,
                                    guru_accuracies=guru_accuracies)

        majority_probability = sum(accs[len(accs) - len(accs) // 2:])
        return majority_probability

    def subtree_accuracy(self, left_accs, right_accs, curr_weight, curr_idx, guru_weights, guru_accuracies):
        """
        Calculate the chance that this guru, along with the gurus underneath it on the guru tree, will result_accuracy in exactly
        i correct votes.
        In left_accs, right_accs, and the returned list, index i is the probability that the guru creating that list
        will result_accuracy in exactly i correct votes. Sum for i > n/2 to get the probability of a majority.
        :param left_accs: list of accuracies of left child; [0] if no left child; None if not yet set
        :param right_accs: list of accuracies of right child or empty list if no right child
        :param curr_weight: weight of the guru currently adding to this
        :param curr_idx: index/label of the guru currently calculating new_accuracies
        :param guru_weights: dictionary mapping each guru index to it's weight
        :param guru_accuracies: dictionary mapping each guru index to it's accuracy
        :return:
        """

        # base cases
        curr_successors = list(self.guru_tree.successors(curr_idx))
        if left_accs is None and len(curr_successors) > 0:
            # check if left subtree exists
                # lowest-valued (or only) child/successor of the current node
            ls = min(curr_successors)
            left_accs = self.subtree_accuracy(left_accs=None,
                                  right_accs=None,
                                  curr_weight=guru_weights[ls],
                                  curr_idx=ls,
                                  guru_weights=guru_weights,
                                  guru_accuracies=guru_accuracies)
        elif left_accs is None:
            left_accs = [1]     # base case: if no children, they have probability 1 of exactly 0 correct votes
        if right_accs is None and len(curr_successors) == 2:
            # check if right subtree exists
                # highest-valued child/successor of the current node, only if two children exist
            rs = max(curr_successors)
            right_accs = self.subtree_accuracy(left_accs=None,
                                  right_accs=None,
                                  curr_weight=guru_weights[rs],
                                  curr_idx=rs,
                                  guru_weights=guru_weights,
                                  guru_accuracies=guru_accuracies)
        elif right_accs is None:
            right_accs = [1]

        max_child_weight = max(len(left_accs)-1, 0) + max(len(right_accs)-1, 0)
        child_accuracies = [0 for _ in range(max_child_weight+curr_weight+1)]   # weight of both subtrees and this plus 1 for 0 weight

        # merge the accuracy of both child subtrees
        # TODO: Doesn't account for one of the subtrees being empty (can we make it [0]?)
        for l_idx, l_acc in enumerate(left_accs):
            for r_idx, r_acc in enumerate(right_accs):

                child_accuracies[l_idx + r_idx] += l_acc * r_acc

        # add this guru's accuracy to that of the combined children's accuracies
        curr_acc = guru_accuracies[curr_idx]
        new_accuracies = child_accuracies.copy()
        for w in range(0, len(child_accuracies)):

            if w >= curr_weight:
                # chance of being incorrect and children getting w correct + chance of children getting w-curr_weight correct
                new_accuracies[w] = child_accuracies[w]*(1-curr_acc) + child_accuracies[w-curr_weight]*curr_acc
            else:
                # weights below curr_weight can only be reached if this guru is incorrect
                new_accuracies[w] = child_accuracies[w] * (1 - curr_acc)

        return new_accuracies


def delegation_forest_from_delegations(delegations):
    """
    Construct a delegation forest from the given list of delegations
    :param delegations: List where value at index i indicates where voter i delegates
    :return:
    """
    f = forest(n=len(delegations))
    for delegator, recipient in enumerate(delegations):
        f.add_delegation(delegator, recipient)

    return f
