from collections import defaultdict
from causaldag.utils import core_utils
import itertools as itr
import numpy as np


class CycleError(Exception):
    def __init__(self):
        super().__init__()
        # def __init__(self, source, target):
        #     self.source = source
        #     self.target = target
        #     message = '%s -> %s will cause a cycle' % (source, target)
        #     super().__init__(message)


class SpouseError(Exception):
    def __init__(self):
        super().__init__()
        # def __init__(self, ancestor, desc):
        #     self.ancestor = ancestor
        #     self.desc = desc
        #     message = '%s <-> %s cannot be added since %s is an ancestor of %s' % (ancestor, desc, ancestor, desc)
        #     super().__init__(message)


class AdjacentError(Exception):
    def __init__(self, node1, node2, arrow_type):
        self.node1 = node1
        self.node2 = node2
        self.arrow_type = arrow_type
        message = '%s %s %s cannot be added since %s and %s are already adjacent' % (
        node1, arrow_type, node2, node1, node2)
        super().__init__(message)


class NeighborError(Exception):
    def __init__(self, node, neighbors=None, parents=None, spouses=None):
        self.node = node
        self.neighbors = neighbors
        self.parents = parents
        self.spouses = spouses
        if self.neighbors:
            message = 'The node %s has neighbors %s. Nodes cannot have neighbors and parents/spouses.' % (
            node, ','.join(map(str, neighbors)))
        elif self.parents:
            message = 'The node %s has parents %s. Nodes cannot have neighbors and parents/spouses.' % (
            node, ','.join(map(str, parents)))
        elif self.spouses:
            message = 'The node %s has spouses %s. Nodes cannot have neighbors and parents/spouses.' % (
            node, ','.join(map(str, spouses)))
        super().__init__(message)


def path2str(path):
    return '->'.join(map(str, path))


class AncestralGraph:
    """
    Base class for ancestral graphs, used to represent causal models with latent variables.
    """

    def __init__(self, nodes=set(), directed=set(), bidirected=set(), undirected=set()):
        self._nodes = nodes.copy()
        self._directed = set()
        self._bidirected = set()
        self._undirected = set()

        self._neighbors = defaultdict(set)
        self._spouses = defaultdict(set)
        self._parents = defaultdict(set)
        self._children = defaultdict(set)
        self._adjacent = defaultdict(set)

        for i, j in directed:
            self._add_directed(i, j)
        for i, j in bidirected:
            self._add_bidirected(i, j)
        for i, j in undirected:
            self._add_undirected(i, j)

    def __eq__(self, other):
        if not isinstance(other, AncestralGraph):
            return False

        same_nodes = self._nodes == other._nodes
        same_directed = self._directed == other._directed
        same_bidirected = self._bidirected == other._bidirected
        same_undirected = self._undirected == other._undirected
        return same_nodes and same_directed and same_bidirected and same_undirected

    def copy(self):
        """
        Return a copy of the current ancestral graph.
        """
        return AncestralGraph(self.nodes, self.directed, self.bidirected, self.undirected)

    def __str__(self):
        return 'Directed edges: %s, Bidirected edges: %s, Undirected edges: %s' % (self._directed, self._bidirected, self._undirected)

    def __repr__(self):
        return str(self)

    # === MUTATORS
    def add_node(self, node):
        """Add a node to the ancestral graph.

        Parameters
        ----------
        node:
            a hashable Python object

        See Also
        --------
        add_nodes_from

        Examples
        --------
        >>> g = cd.AncestralGraph()
        >>> g.add_node(1)
        >>> g.add_node(2)
        >>> len(g.nodes)
        2
        """
        self._nodes.add(node)

    def add_nodes_from(self, nodes):
        """Add a node to the ancestral graph.

        Parameters
        ----------
        nodes:
            an iterable of hashable Python objects

        See Also
        --------
        add_node

        Examples
        --------
        >>> g = cd.AncestralGraph()
        >>> g.add_nodes_from({1, 2})
        >>> len(g.nodes)
        2
        """
        for node in nodes:
            self._nodes.add(node)

    def _check_ancestral(self):
        self.topological_sort()

    def _mark_children_visited(self, node, any_visited, curr_path_visited, curr_path, stack):
        any_visited[node] = True
        curr_path_visited[node] = True
        curr_path.append(node)
        for child in self._children[node]:
            if not any_visited[child]:
                self._mark_children_visited(child, any_visited, curr_path_visited, curr_path, stack)
            elif curr_path_visited[child]:
                cycle = curr_path + [child]
                raise CycleError
        for spouse in self._spouses[node]:
            if curr_path_visited[spouse]:
                raise SpouseError
        curr_path.pop()
        curr_path_visited[node] = False
        stack.append(node)

    def topological_sort(self):
        any_visited = {node: False for node in self._nodes}
        curr_path_visited = {node: False for node in self._nodes}
        curr_path = []
        stack = []
        for node in self._nodes:
            if not any_visited[node]:
                self._mark_children_visited(node, any_visited, curr_path_visited, curr_path, stack)
        return list(reversed(stack))

    def add_directed(self, i, j):
        self._add_directed(i, j)
        try:
            self._check_ancestral()
        except CycleError as e:
            self.remove_directed(i, j)
            raise e

    def add_bidirected(self, i, j):
        self._add_bidirected(i, j)
        try:
            self._add_bidirected(i, j)
        except CycleError as e:
            self.remove_bidirected(i, j)
            raise e

    def add_undirected(self, i, j):
        self._add_undirected(i, j)

    def _add_directed(self, i, j, ignore_error=False):
        if self.has_directed(i, j):
            return

        # === CHECK REMAINS ANCESTRAL
        if self._neighbors[j]:
            raise NeighborError(j, self._neighbors[j])

        # === CHECK i AND j NOT ALREADY ADJACENT
        if i in self._adjacent[j]:
            if ignore_error:
                if self.has_directed(j, i):
                    self.remove_directed(j, i)
                elif self.has_bidirected(i, j):
                    self.remove_bidirected(i, j)
                else:
                    self.remove_undirected(i, j)
            else:
                raise AdjacentError(i, j, '->')

        self._nodes.add(i)
        self._nodes.add(j)
        self._directed.add((i, j))
        self._parents[j].add(i)
        self._children[i].add(j)

        self._adjacent[i].add(j)
        self._adjacent[j].add(i)

    def _add_bidirected(self, i, j, ignore_error=False):
        if self.has_bidirected(i, j):
            return

        # === CHECK REMAINS ANCESTRAL
        if self._neighbors[i]:
            raise NeighborError(i, neighbors=self._neighbors[i])
        if self._neighbors[j]:
            raise NeighborError(j, neighbors=self._neighbors[j])

        # === CHECK i AND j NOT ALREADY ADJACENT
        if i in self._adjacent[j]:
            if ignore_error:
                if self.has_directed(i, j):
                    self.remove_directed(i, j)
                elif self.has_directed(j, i):
                    self.remove_directed(j, i)
                else:
                    self.remove_undirected(i, j)
            else:
                raise AdjacentError(i, j, '<->')

        self._nodes.add(i)
        self._nodes.add(j)
        self._bidirected.add(tuple(sorted((i, j))))
        self._spouses[j].add(i)
        self._spouses[i].add(j)

        self._adjacent[i].add(j)
        self._adjacent[j].add(i)

    def _add_undirected(self, i, j, ignore_error=False):
        if self.has_undirected(i, j):
            return

        # === CHECK REMAINS ANCESTRAL
        if self._parents[i]:
            raise NeighborError(i, parents=self._parents[i])
        if self._spouses[i]:
            raise NeighborError(i, spouses=self._spouses[i])
        if self._parents[j]:
            raise NeighborError(j, parents=self._parents[j])
        if self._spouses[j]:
            raise NeighborError(j, spouses=self._spouses[j])

        # === CHECK i AND j NOT ALREADY ADJACENT
        if i in self._adjacent[j]:
            if ignore_error:
                if self.has_directed(i, j):
                    self.remove_directed(i, j)
                elif self.has_directed(j, i):
                    self.remove_directed(j, i)
                else:
                    self.remove_bidirected(i, j)
            else:
                raise AdjacentError(i, j, '-')

        self._nodes.add(i)
        self._nodes.add(j)
        self._undirected.add(tuple(sorted((i, j))))
        self._neighbors[j].add(i)
        self._neighbors[i].add(j)

        self._adjacent[i].add(j)
        self._adjacent[j].add(i)

    def remove_node(self, node, ignore_error=False):
        try:
            self._nodes.remove(node)
            for parent in self._parents[node]:
                self._children[parent].remove(node)
                self._adjacent[parent].remove(node)
                self._directed.remove((parent, node))
            for child in self._children[node]:
                self._parents[child].remove(node)
                self._adjacent[child].remove(node)
                self._directed.remove((node, child))
            for spouse in self._spouses[node]:
                self._spouses[spouse].remove(node)
                self._adjacent[spouse].remove(node)
                self._bidirected.remove(tuple(sorted((spouse, node))))
            for nbr in self._neighbors[node]:
                self._neighbors[nbr].remove(node)
                self._adjacent[nbr].remove(node)
                self._undirected.remove(tuple(sorted((nbr, node))))

            del self._children[node]
            del self._parents[node]
            del self._spouses[node]
            del self._neighbors[node]
            del self._adjacent[node]
        except KeyError as e:
            if ignore_error:
                pass
            else:
                raise e

    def remove_directed(self, i, j, ignore_error=False):
        try:
            self._directed.remove((i, j))
            self._children[i].remove(j)
            self._parents[j].remove(i)
            self._adjacent[i].remove(j)
            self._adjacent[j].remove(i)
        except KeyError as e:
            if ignore_error:
                pass
            else:
                raise e

    def remove_bidirected(self, i, j, ignore_error=False):
        try:
            self._bidirected.remove(tuple(sorted((i, j))))
            self._spouses[i].remove(j)
            self._spouses[j].remove(i)
            self._adjacent[i].remove(j)
            self._adjacent[j].remove(i)
        except KeyError as e:
            if ignore_error:
                pass
            else:
                raise e

    def remove_undirected(self, i, j, ignore_error=False):
        try:
            self._undirected.remove(tuple(sorted((i, j))))
            self._neighbors[i].remove(j)
            self._neighbors[j].remove(i)
            self._adjacent[i].remove(j)
            self._adjacent[j].remove(i)
        except KeyError as e:
            if ignore_error:
                pass
            else:
                raise e

    # === PROPERTIES
    @property
    def nodes(self):
        return self._nodes.copy()

    @property
    def nnodes(self):
        return len(self._nodes)

    @property
    def directed(self):
        return self._directed.copy()

    @property
    def num_directed(self):
        return len(self._directed)

    @property
    def bidirected(self):
        return self._bidirected.copy()

    @property
    def num_bidirected(self):
        return len(self._bidirected)

    @property
    def undirected(self):
        return self._undirected.copy()

    @property
    def num_undirected(self):
        return len(self._undirected)

    @property
    def num_edges(self):
        return self.num_directed + self.num_bidirected + self.num_undirected

    @property
    def skeleton(self):
        return {frozenset({i, j}) for i, j in self._bidirected | self._undirected | self._directed}

    def children_of(self, i):
        return self._children[i].copy()

    def parents_of(self, i):
        return self._parents[i].copy()

    def spouses_of(self, i):
        return self._spouses[i].copy()

    def neighbors_of(self, i):
        return self._neighbors[i].copy()

    def _add_ancestors(self, ancestors, node, exclude_arcs=set()):
        for parent in self._parents[node]:
            if parent not in ancestors and (parent, node) not in exclude_arcs:
                ancestors.add(parent)
                self._add_ancestors(ancestors, parent, exclude_arcs=exclude_arcs)

    def ancestors_of(self, node, exclude_arcs=set()):
        """Return the nodes upstream of node

        Parameters
        ----------
        node:
            The node.

        See Also
        --------
        descendants_of

        Return
        ------
        Set[node]
            Return all nodes j such that there is a directed path from j to node.

        Example
        -------
        """
        ancestors = set()
        self._add_ancestors(ancestors, node, exclude_arcs=exclude_arcs)
        return ancestors

    def descendants_of(self, i):
        raise NotImplementedError

    def has_directed(self, i, j):
        return (i, j) in self._directed

    def has_bidirected(self, i, j):
        return tuple(sorted((i, j))) in self._bidirected

    def has_undirected(self, i, j):
        return tuple(sorted((i, j))) in self._undirected

    def has_any_edge(self, i, j):
        return self.has_directed(i, j) or self.has_directed(j, i) or self.has_bidirected(i, j) or self.has_undirected(i, j)

    def vstructures(self):
        vstructs = set()
        for node in self._nodes:
            for p1, p2 in itr.combinations(self._parents[node] | self._spouses[node], 2):
                if not self.has_any_edge(p1, p2):
                    p1_, p2_ = sorted((p1, p2))
                    vstructs.add((p1_, node, p2))
        return vstructs

    def colliders(self):
        return {node for node in self._nodes if len(self._parents[node] | self._spouses[node]) >= 2}

    def discriminating_paths(self):
        colliders = self.colliders()
        discriminating_paths = {}
        for j, parents in self._parents.items():  # potential endpoints of discriminating paths
            if not parents:
                break
            nonadjacent = self._nodes - parents - self._children[j] - self._spouses[j] - {j}
            for i in nonadjacent:  # potential start points of discriminating paths
                # search all paths that satisfy discriminating path criteria
                path_queue = [
                    [i, k]
                    for k in self._spouses[i] | self._children[i]
                    if k in colliders and j in self._children[k]
                ]
                while path_queue:
                    path = path_queue.pop(0)
                    final_node = path[-1]
                    # check if path is discriminating for the next node
                    for k in filter(lambda k: k not in path, self._spouses[final_node]):
                        if j in self._spouses[k]:
                            full_path = path.copy()
                            full_path.extend([k, j])
                            discriminating_paths[tuple(full_path)] = 'c'
                        elif j in self._children[k]:
                            full_path = path.copy()
                            full_path.extend([k, j])
                            discriminating_paths[tuple(full_path)] = 'n'
                    for k in filter(lambda k: k not in path, self._parents[final_node]):
                        if j in self._children[k]:
                            full_path = path.copy()
                            full_path.extend([k, j])
                            discriminating_paths[tuple(full_path)] = 'n'

                    # extend path
                    for k in self._spouses[final_node]:
                        if k not in path and k in colliders and j in self._children[k]:
                            new_path = path.copy()
                            new_path.append(k)
                            path_queue.append(new_path)
        return discriminating_paths

    # === ???
    def to_maximal(self):
        for i, j in itr.combinations(self._nodes, r=2):
            if not self.has_any_edge(i, j):
                never_msep = not any(self.msep(i, j, S) for S in core_utils.powerset(self._nodes - {i, j}))
                if never_msep: self.add_bidirected(i, j)

    def to_pag(self):
        raise NotImplementedError

    # === CONVERTERS
    def to_amat(self):
        amat = np.zeros([self.nnodes, self.nnodes])
        for i, j in self.directed:
            amat[i, j] = 2
            amat[j, i] = 3
        for i, j in self.bidirected:
            amat[i, j] = 2
            amat[j, i] = 2
        for i, j in self.undirected:
            amat[i, j] = 3
            amat[j, i] = 3
        return amat

    def from_amat(self, amat):
        p = amat.shape[0]
        directed = set()
        bidirected = set()
        undirected = set()
        for i, j in itr.combinations(set(range(p)), 2):
            vij = amat[i, j]
            vji = amat[j, i]
            if vij == 2 and vji == 3:
                directed.add((i, j))
            if vij == 3 and vji == 2:
                directed.add((j, i))
            if vij == 2 and vji == 2:
                bidirected.add((i, j))
            if vij == 3 and vji == 3:
                undirected.add((i, j))
        return AncestralGraph(set(range(p)), directed, bidirected, undirected)

    # === COMPARISON
    def markov_equivalent(self, other):
        same_skeleton = self.skeleton == other.skeleton
        same_vstructures = self.vstructures() == other.vstructures()

        self_discriminating_paths = self.discriminating_paths()
        other_discriminating_paths = other.discriminating_paths()
        shared_disc_paths = set(self_discriminating_paths.keys()) & set(other_discriminating_paths)
        same_discriminating = all(
            self_discriminating_paths[path] == other_discriminating_paths[path]
            for path in shared_disc_paths
        )

        return same_skeleton and same_vstructures and same_discriminating

    def shd_skeleton(self, other):
        return len(self.skeleton.symmetric_difference(other.skeleton))

    def as_hashed(self):
        return frozenset(self._directed), frozenset(self._bidirected), frozenset(self._undirected)

    # === Algorithms
    def _add_upstream(self, upstream, node):
        for parent in self._parents[node]:
            if parent not in upstream:
                upstream.add(parent)
                self._add_upstream(upstream, parent)

    def _is_collider(self, u, v, w):
        """return True if u-v-w is a collider"""
        if v in self._children[u] and v in self._children[w]:
            return True
        elif v in self._children[u] and v in self._spouses[w]:
            return True
        elif v in self._spouses[u] and v in self._children[w]:
            return True
        elif v in self._spouses[u] and v in self._spouses[w]:
            return True
        else:
            return False

    def legitimate_mark_changes(self, verbose=False):
        """
        Return directed edges that can be changed to bidirected edges, and bidirected edges that can be changed to
        directed edges.

        Return
        ------
        (mark_changes_dir, mark_changes_bidir)
            Directed edges that can be changed to bidirected edges, and bidirected edges that can be changed to directed
            edges (which will be the new directed edge).

        Example
        -------
        >>> g = cd.AncestralGraph(directed={(0, 1)}, bidirected={(1, 2)})
        >>> g.legitimate_mark_changes()
        ({(0, 1)}, {(2, 1)})
        """
        if self._undirected:
            raise ValueError('Only defined for DMAGs')

        disc_paths = self.discriminating_paths()

        mark_changes_dir = set()
        for i, j in self._directed:
            if verbose: print(f'{i}->{j} => {i}<->{j} ?')
            # FIRST CONDITIONS
            parents_condition = all(self.has_directed(parent, j) for parent in self._parents[i])
            if not parents_condition:
                if verbose: print('Failed parents condition')
                continue
            spouses_condition = all(self.has_any_edge(spouse, j) for spouse in self._spouses[i])
            if not spouses_condition:
                if verbose: print('Failed spouses condition')
                continue

            # SECOND CONDITION
            disc_paths_for_i = [path for path in disc_paths.keys() if path[-2] == i]
            disc_paths_condition = all(path[-1] != j for path in disc_paths_for_i)
            if not disc_paths_condition:
                if verbose: print('Failed discriminating path condition')
                continue

            # FINAL CONDITION
            if i in self.ancestors_of(j, exclude_arcs={(i, j)}):
                if verbose: print('Failed ancestral condition')
                continue

            if verbose: print('Passed')
            mark_changes_dir.add((i, j))

        mark_changes_bidir = set()
        for i, j in self._bidirected | set(map(reversed, self._bidirected)):
            if verbose: print(f'{i}<->{j} => {i}->{j} ?')
            # FIRST CONDITIONS
            parents_condition = all(self.has_directed(parent, j) for parent in self._parents[i])
            if not parents_condition:
                if verbose: print('Failed parents condition')
                continue
            spouses_condition = all(self.has_any_edge(spouse, j) for spouse in self._spouses[i] if spouse != j)
            if not spouses_condition:
                if verbose: print('Failed spouses condition')
                continue

            # SECOND CONDITION
            disc_paths_for_i = [path for path in disc_paths.keys() if path[-2] == i]
            disc_paths_condition = all(path[-1] != j for path in disc_paths_for_i)
            if not disc_paths_condition:
                if verbose: print('Failed discriminating path condition')
                continue

            # FINAL CONDITION
            if i in self.ancestors_of(j):
                if verbose: print('Failed ancestral condition')
                continue

            if verbose: print('Passed')
            mark_changes_bidir.add((i, j))

        return mark_changes_dir, mark_changes_bidir

    def msep(self, A, B, C=set()):
        """
        Check whether A and B are m-separated given C, using the Bayes ball algorithm.


        """
        # type coercion
        A = core_utils.to_set(A)
        B = core_utils.to_set(B)
        C = core_utils.to_set(C)

        # shade ancestors of C
        shaded_nodes = set(C)
        for node in C:
            self._add_upstream(shaded_nodes, node)

        visited = set()
        # marks whether the node has been encountered along a path where it has a tail or an arrowhead
        _t = 'tail'  # tail
        _a = 'arrowhead'  # arrowhead

        schedule = {(node, _t) for node in A}
        while schedule:
            node, _dir = schedule.pop()
            if node in B: return False
            if (node, _dir) in visited: continue
            visited.add((node, _dir))
            # print(node, _dir)

            # if coming through a tail, won't encounter v-structure
            if _dir == _t and node not in C:
                schedule.update({(parent, _t) for parent in self._parents[node]})
                schedule.update({(child, _a) for child in self._children[node]})
                schedule.update({(spouse, _a) for spouse in self._spouses[node]})
                schedule.update({(nbr, _t) for nbr in self._neighbors[node]})

            if _dir == _a:
                # if coming through an arrowhead and see shaded node, can go through v-structure
                if node in shaded_nodes:
                    schedule.update({(parent, _t) for parent in self._parents[node]})
                    schedule.update({(spouse, _a) for spouse in self._spouses[node]})

                # if coming through an arrowhead and see unconditioned node, can go through children and neighbors
                if node not in C:
                    schedule.update({(child, _a) for child in self._children[node]})
                    schedule.update({(nbr, _a) for nbr in self._neighbors[node]})

        return True

    def msep_from_given(self, A, C=set()):
        """Find all nodes m-seperated from A given C using algorithm similar to that in Geiger, D., Verma, T., & Pearl, J. (1990). Identifying independence in Bayesian networks. Networks, 20(5), 507-534."""

        A = core_utils.to_set(A)
        C = core_utils.to_set(C)

        determined = set()
        descendants = set()

        for c in C:
            determined.add(c)
            descendants.add(c)
            self._add_upstream(descendants, c)

        reachable = set()
        i_links = set()
        labeled_links = set()

        for a in A:
            i_links.add((None, a))
            reachable.add(a)

        while True:
            i_p_1_links = set()
            # Find all unlabled links v->w adjacent to at least one link u->v labeled i, such that (u->v,v->w) is a legal pair.
            for link in i_links:
                u, v = link
                for w in self._adjacent[v]:
                    if not u == w and (v, w) not in labeled_links:
                        if self._is_collider(u, v, w):  # Is collider?
                            if v in descendants:
                                i_p_1_links.add((v, w))
                                reachable.add(w)
                        else:  # Not collider
                            if v not in determined:
                                i_p_1_links.add((v, w))
                                reachable.add(w)

            if len(i_p_1_links) == 0:
                break

            labeled_links = labeled_links.union(i_links)
            i_links = i_p_1_links

        return self._nodes.difference(A).difference(C).difference(reachable)


if __name__ == '__main__':
    g = AncestralGraph(nodes=set(range(1, 5)), directed={(1, 2), (2, 4), (3, 2), (3, 4)})
    disc_paths = g.discriminating_paths()



