from collections import defaultdict

from beliefs.types.Node import Node
from beliefs.factors.BernoulliOrCPD import BernoulliOrCPD


class EdgesHelper:
    """Class with convenience methods for working with edges."""
    def __init__(self, edges):
        self.edges = edges

    def get_label_to_children_dict(self):
        """returns dictionary keyed on label, with value a set of children"""
        label_to_children_dict = defaultdict(set)
        for parent, child in self.edges:
            label_to_children_dict[parent].add(child)
        return label_to_children_dict

    def get_label_to_parents_dict(self):
        """returns dictionary keyed on label, with value a set of parents
        Only used to help create dummy factors from edges (not for algo).
        """
        label_to_parents_dict = defaultdict(set)

        for parent, child in self.edges:
            label_to_parents_dict[child].add(parent)
        return label_to_parents_dict

    def get_labels_from_edges(self):
        """Return the set of labels that comprise the vertices of a list of edge tuples."""
        all_labels = set()
        for parent, child in self.edges:
            all_labels.update({parent, child})
        return all_labels

    def create_cpds_from_edges(self, CPD=BernoulliOrCPD):
        """
        Create factors from list of edges.

        Input:
          cpd: a factor class, assumed initialization takes in a label_id, the label_id of
               the child (should = label_id of the node), and set of label_ids of parents.

        Returns:
          factors: a set of (unique) factors of the graph
        """
        labels = self.get_labels_from_edges()
        label_to_parents = self.get_label_to_parents_dict()

        factors = set()

        for label in labels:
            parents = label_to_parents[label]
            cpd = CPD(label, parents)
            factors.add(cpd)
        return factors

    def get_label_to_factor_dict(self, CPD=BernoulliOrCPD):
        """Create a dictionary mapping each label_id to the CPD/factor where
        that label_id is a child.

        Returns:
          label_to_factor: dict mapping each label to the cpd that
                          has that label as a child.
        """
        factors = self.create_cpds_from_edges(CPD=CPD)

        label_to_factor = dict()
        for factor in factors:
            label_to_factor[factor.child] = factor
        return label_to_factor

    def get_label_to_node_dict(self, CPD=BernoulliOrCPD):
        """Create a dictionary mapping each label_id to a Node instance.

        Returns:
          label_to_node: dict mapping each label to the node that has that
                         label as a label_id.
        """
        nodes = self.create_nodes_from_edges()

        label_to_node = dict()
        for node in nodes:
            label_to_node[node.label_id] = node
        return label_to_node

    def get_label_to_node_dict_for_manual_cpds(self, cpds_list):
        """Create a dictionary mapping each label_id to a node that is
        instantiated with a manually defined pgmpy factor instance.

        Input:
          cpds_list - list of instances of pgmpy factors, e.g. TabularCPD

        Returns:
          label_to_node: dict mapping each label to the node that has that
                         label as a label_id.
        """
        label_to_children = self.get_label_to_children_dict()
        label_to_parents = self.get_label_to_parents_dict()

        label_to_node = dict()
        for cpd in cpds_list:
            label_id = cpd.variable

            node = Node(label_id=label_id,
                        children=label_to_children[label_id],
                        parents=label_to_parents[label_id],
                        cardinality=2,
                        cpd=cpd)
            label_to_node[label_id] = node

        return label_to_node

    def create_nodes_from_edges(self, node_class):
        """
        Create instances of the node_class.  Assumes the node class is
        initialized by label_id, children, and parents.

        Returns:
          nodes: a set of (unique) nodes of the graph
        """
        labels = self.get_labels_from_edges()
        labels_to_parents = self.get_label_to_parents_dict()
        labels_to_children = self.get_label_to_children_dict()

        nodes = set()

        for label in labels:
            parents = list(labels_to_parents[label])
            children = list(labels_to_children[label])

            node = node_class(label_id=label,
                              children=children,
                              parents=parents)
            nodes.add(node)
        return nodes