"""
Simple Dependency Graph (SDG)

This module supports Simple Dependency Graphs (SDGs) which are a general
form of graph that has at most one edge between any pair of nodes.

The two endpoints of each edge can be a head ">", tail "-" or circle "o"
(which means either head or tail).

This format can represent most dependency graph types including:
 - Markov Graphs
 - Directed Acyclic Graphs (DAGs)
 - Partially Directed Acyclic Graphs (PDAGs)
 - Maximal Ancestral Graphs (MAGs)
 - Partial Ancestral Graphs (PAGs)
"""

from typing import Dict, List, Optional, Set, Tuple, Union

from numpy import zeros
from pandas import DataFrame


class SDG:
    """Base class for simple dependency graphs (one edge between vertices).

    Simple Dependency Graphs (SDGs) are a general form of graph that has at
    most one edge between any pair of nodes. The two endpoints of each edge
    can be a     head ">", tail "-" or circle "o" (which means either
    head or tail).

    This format can represent most dependency graph types including:
    - Markov Graphs
    - Directed Acyclic Graphs (DAGs)
    - Partially Directed Acyclic Graphs (PDAGs)
    - Maximal Ancestral Graphs (MAGs)
    - Partial Ancestral Graphs (PAGs)

    Args:
        nodes: Nodes present in the graph.
        edges: Edges which define the graph connections as list of tuples:
            (node1, dependency symbol, node2).

    Attributes:
        nodes: Graph nodes in alphabetical order.
        edges: Graph edges {(node1, node2): EdgeType}.
        is_directed: Graph only has directed (causal) edges.
        is_partially_directed: Graph is partially directed.
        parents: Parents of node {node: [parents]}.
        has_directed_cycles: Contains any directed cycles.

    Raises:
        TypeError: If nodes and edges are not both lists.
        ValueError: If node or edge is invalid.
    """

    def __init__(
        self, nodes: List[str], edges: List[Tuple[str, str, str]]
    ) -> None:
        """Initialize a Simple Dependency Graph.

        Args:
            nodes: List of node names in the graph.
            edges: List of edge tuples in format
                (node1, edge_type_symbol, node2).
        """
        # Import EdgeType locally to avoid circular imports
        from causaliq_core.graph import EdgeType

        if type(nodes) is not list or type(edges) is not list:
            raise TypeError("graph edges and nodes not both lists")

        # Validate nodes specified

        self.nodes: List[str] = []
        for node in nodes:
            if type(node) is not str:
                raise TypeError("graph has non-string node name")
            if not node:
                raise ValueError("graph has empty node name")
            if node in self.nodes:
                raise ValueError("graph has duplicate node names")
            self.nodes.append(node)
        self.nodes.sort()

        # Validate edges specified

        self.edges: Dict[Tuple[str, str], EdgeType] = {}
        edge_types = {e.value[3]: e for e in EdgeType}
        self.is_directed: bool = True
        self.is_partially_directed: bool = True
        self.parents: Dict[str, List[str]] = {}
        self.has_directed_cycles: bool = False

        for edge in edges:
            if type(edge) is not tuple or len(edge) != 3:
                raise TypeError("graph has non-triple edge")
            for part in edge:
                if type(part) is not str:
                    raise TypeError("graph has non-string edge triple")
            if edge[0] == edge[2]:
                raise ValueError("graph has cyclic edge")
            if edge[1] not in edge_types:
                raise TypeError(
                    'graph edge has unknown type symbol: "{}"'.format(edge[1])
                )
            if edge[0] not in self.nodes or edge[2] not in self.nodes:
                raise ValueError("graph edge has unknown node")
            if (edge[0], edge[2]) in self.edges or (
                edge[2],
                edge[0],
            ) in self.edges:
                raise ValueError("graph has duplicate edges")

            # Check if graph directed or partially directed
            edge_type = edge_types[edge[1]]
            if edge_type != EdgeType.DIRECTED:
                self.is_directed = False
                if edge_type != EdgeType.UNDIRECTED:
                    self.is_partially_directed = False

            if edge_type == EdgeType.DIRECTED:

                # collect parents for directional links

                if edge[2] in self.parents:
                    self.parents[edge[2]].append(edge[0])
                else:
                    self.parents[edge[2]] = [edge[0]]

            elif (
                edge_type
                in (
                    EdgeType.UNDIRECTED,
                    EdgeType.BIDIRECTED,
                    EdgeType.NONDIRECTED,
                )
                and edge[0] > edge[2]
            ):

                # ensure nodes in alphabetical order for non-directional edges

                edge = (edge[2], edge[1], edge[0])

            # parents held in alphabetical order

            self.parents = {k: sorted(v) for k, v in self.parents.items()}

            # store the edge in dictionary {(node1, node2): edge_type}

            self.edges[(edge[0], edge[2])] = edge_type

            # see whether has directed cycles

            self.has_directed_cycles = (
                True
                if self.partial_order(self.parents, self.nodes) is None
                else False
            )

    def rename(self, name_map: Dict[str, str]) -> None:
        """Rename nodes in place according to name map.

        Args:
            name_map: Name mapping {name: new name}. Must have mapping for
                every node.

        Raises:
            TypeError: With bad arg type.
            ValueError: With bad arg values e.g. unknown node names.
        """
        if not isinstance(name_map, dict) or not all(
            [
                isinstance(k, str) and isinstance(v, str)
                for k, v in name_map.items()
            ]
        ):
            raise TypeError("SDG.rename() bad arg types")

        if set(name_map.keys()) != set(self.nodes):
            raise ValueError("SDG.raname() bad arg values")

        # Change node names and node names in edges

        nodes = [name_map[n] if n in name_map else n for n in self.nodes]
        edges = [
            (
                (name_map[e[0]] if e[0] in name_map else e[0]),
                t.value[3],
                (name_map[e[1]] if e[1] in name_map else e[1]),
            )
            for e, t in self.edges.items()
        ]

        # Re-instantiate object with new names
        new_sdg = SDG(nodes, edges)
        self.nodes = new_sdg.nodes
        self.edges = new_sdg.edges
        self.is_directed = new_sdg.is_directed
        self.is_partially_directed = new_sdg.is_partially_directed
        self.parents = new_sdg.parents
        self.has_directed_cycles = new_sdg.has_directed_cycles

    @classmethod
    def partial_order(
        cls,
        parents: Dict[str, List[str]],
        nodes: Optional[Union[List[str], Set[str]]] = None,
        new_arc: Optional[Tuple[str, str]] = None,
    ) -> Optional[List[Set[str]]]:
        """Return partial topological ordering for the directed part of a
        graph.

        The graph is specified by list of parents for each node.

        Args:
            parents: Parents of each node {node: [parents]}.
            nodes: Optional complete list of nodes including parentless ones
                for use if parents argument doesn't include them already.
            new_arc: A new arc (n1, n2) to be added before order is evaluated.
                If the opposing arc is implied in parents then it is removed so
                that arc reversal is also supported. This argument facilitates
                seeing whether an arc addition or reversal would create a
                cycle.

        Returns:
            Nodes in a partial topological order as list of sets or None if
            there is no ordering which means the graph is cyclic.
        """
        if not isinstance(parents, dict) or (
            nodes is not None
            and not isinstance(nodes, list)
            and not isinstance(nodes, set)
        ):
            raise TypeError("DAG.partial_order bad arg type")

        parents_copy = {n: set(p) for n, p in parents.items()}
        if nodes is not None:
            parents_copy.update({n: set() for n in nodes if n not in parents})
        if new_arc is not None:
            parents_copy[new_arc[0]] = parents_copy[new_arc[0]] - {new_arc[1]}
            parents_copy[new_arc[1]] = parents_copy[new_arc[1]] | {new_arc[0]}

        order = []
        while len(parents_copy):
            roots = {n for n in parents_copy if not parents_copy[n]}
            if not len(roots):
                return None
            order.append(roots)
            for root in roots:
                parents_copy.pop(root)
            for node in list(parents_copy.keys()):
                parents_copy[node].difference_update(roots)

        return order

    def is_DAG(self) -> bool:
        """Return whether graph is a Directed Acyclic Graph (DAG).

        Returns:
            True if graph is a DAG, False otherwise.
        """
        return self.is_directed and self.has_directed_cycles is False

    def is_PDAG(self) -> bool:
        """Return whether graph is a Partially Directed Acyclic Graph (PDAG).

        Returns:
            True if graph is a PDAG, False otherwise.
        """
        from causaliq_core.graph import EdgeType

        if self.is_directed:
            return self.is_DAG()

        if not self.is_partially_directed:
            return False

        # Check if there are any cycles in directed part of graph

        arcs = [
            (a[0], "->", a[1])
            for a, t in self.edges.items()
            if t == EdgeType.DIRECTED
        ]
        return not SDG(self.nodes, arcs).has_directed_cycles

    def undirected_trees(
        self,
    ) -> List[Set[Union[Tuple[str, str], Tuple[str, None]]]]:
        """Return undirected trees present in graph.

        Returns:
            List of trees, each tree a set of tuples representing edges in tree
            (n1, n2) or a single isolated node (n1, None).
        """
        trees = []
        edges = {
            (e[0], e[1]) if e[1] > e[0] else (e[1], e[0])
            for e in self.edges.keys()
        }
        isolated = set(self.nodes)
        while len(edges):
            # Start with the lexicographically smallest edge for determinism
            first_edge = min(edges)
            tree: Set[Union[Tuple[str, str], Tuple[str, None]]] = {first_edge}
            t_nodes = {first_edge[0], first_edge[1]}

            # Keep adding connected edges until no more can be added
            growing = True
            while growing:
                growing = False
                for edge in sorted(edges):  # Sort for deterministic iteration
                    if edge not in tree and (
                        (edge[0] in t_nodes and edge[1] not in t_nodes)
                        or (edge[1] in t_nodes and edge[0] not in t_nodes)
                    ):
                        t_nodes = t_nodes | {edge[0], edge[1]}
                        tree = tree | {edge}
                        growing = True

            # Remove only the actual edges (not isolated nodes) from edges set
            tree_edges = {edge for edge in tree if edge[1] is not None}
            edges = edges - tree_edges
            isolated = isolated - t_nodes
            trees.append(tree)
        trees.extend(
            [{(n, None)} for n in sorted(isolated)]
        )  # Sort for determinism
        return sorted(trees, key=lambda t: min(t))

    def components(self) -> List[List[str]]:
        """Return components present in graph.

        Uses tree search algorithm to span the undirected graph to identify
        nodes in individual trees which are the spanning tree of each
        component.

        Returns:
            List of lists, each a list of sorted nodes in component.
        """
        components = []
        edges = {
            (e[0], e[1]) if e[1] > e[0] else (e[1], e[0])
            for e in self.edges.keys()
        }  # skeleton of graph
        nodes = set(self.nodes)  # nodes yet to be included in components

        while len(nodes):
            c_nodes = {list(nodes)[0]}  # put 1st arbitrary node in component
            growing = True
            while growing:
                growing = False
                for edge in edges:  # look for edges across component boundary
                    if (edge[0] in c_nodes and edge[1] not in c_nodes) or (
                        edge[1] in c_nodes and edge[0] not in c_nodes
                    ):
                        c_nodes = c_nodes | {edge[0], edge[1]}
                        growing = True

            # component identified, remove its nodes from consideration

            nodes = nodes - c_nodes
            components.append(sorted(c_nodes))

        return sorted(components)

    def number_components(self) -> int:
        """Return number of components (including unconnected nodes) in graph.

        Returns:
            Number of components.
        """
        return len(self.components())

    def to_adjmat(self) -> DataFrame:
        """Return an adjacency matrix representation of the graph.

        Returns:
            Adjacency matrix as a pandas DataFrame.
        """
        size = len(self.nodes)
        adjmat: DataFrame = DataFrame(
            zeros(shape=(size, size), dtype="int8"), columns=self.nodes
        )
        adjmat[""] = self.nodes
        adjmat = adjmat.set_index("")
        for nodes, type in self.edges.items():
            # print('{}: {}'.format(nodes, type.value[0]))
            adjmat.loc[nodes[0], nodes[1]] = type.value[0]

        return adjmat

    def __str__(self) -> str:
        """Return a human-readable description of the graph.

        Returns:
            String description of graph.
        """
        if not self.nodes:
            return "Empty graph"

        number_components = self.number_components()
        kind = (
            "DAG"
            if self.is_directed
            else ("PDAG" if self.is_partially_directed else "SDG")
        )
        desc = (
            "{} ({}), {} node".format(kind, type(self), len(self.nodes))
            + ("" if len(self.nodes) == 1 else "s")
            + ", {} edge".format(len(self.edges))
            + ("" if len(self.edges) == 1 else "s")
            + " and {} component".format(number_components)
            + ("" if number_components == 1 else "s")
        )

        for node in self.nodes:
            children = [
                t.value[3] + str(e[1])
                for e, t in self.edges.items()
                if e[0] == node
            ]
            desc += "\n{}: {}".format(node, " ".join(children))
        return desc

    def __eq__(self, other: object) -> bool:
        """Test if graph is identical to this one.

        Args:
            other: Graph to compare with self.

        Returns:
            True if other is identical to self.
        """
        return (
            isinstance(other, SDG)
            and self.nodes == other.nodes
            and self.edges == other.edges
        )

    def compress(self) -> bytes:
        """Compress graph to compact binary representation.

        Format:

        - 2 bytes: number of nodes (uint16, big-endian)
        - For each node: 2 bytes name length + UTF-8 encoded name
        - 2 bytes: number of edges (uint16, big-endian)
        - For each edge: 2 bytes source index + 2 bytes target index +
          1 byte edge type

        Returns:
            Compact binary representation of the graph.

        Raises:
            ValueError: If graph has more than 65535 nodes or edges.
        """
        if len(self.nodes) > 65535:
            raise ValueError("SDG.compress() graph has too many nodes")
        if len(self.edges) > 65535:
            raise ValueError("SDG.compress() graph has too many edges")

        # Build node index for edge compression
        node_index = {node: i for i, node in enumerate(self.nodes)}

        # Compress number of nodes
        data = len(self.nodes).to_bytes(2, "big")

        # Compress each node name
        for node in self.nodes:
            name_bytes = node.encode("utf-8")
            if len(name_bytes) > 65535:
                raise ValueError("SDG.compress() node name too long")
            data += len(name_bytes).to_bytes(2, "big")
            data += name_bytes

        # Compress number of edges
        data += len(self.edges).to_bytes(2, "big")

        # Compress each edge (source index, target index, edge type)
        for (source, target), edge_type in self.edges.items():
            data += node_index[source].to_bytes(2, "big")
            data += node_index[target].to_bytes(2, "big")
            data += edge_type.value[0].to_bytes(1, "big")

        return data

    @classmethod
    def decompress(cls, data: bytes) -> "SDG":
        """Decompress graph from compact binary representation.

        Args:
            data: Binary data from SDG.compress().

        Returns:
            Reconstructed SDG instance.

        Raises:
            TypeError: If data is not bytes.
            ValueError: If data is invalid or corrupted.
        """
        from causaliq_core.graph import EdgeType

        if not isinstance(data, bytes):
            raise TypeError("SDG.decompress() data must be bytes")

        if len(data) < 2:
            raise ValueError("SDG.decompress() data too short")

        pos = 0

        # Decompress number of nodes
        num_nodes = int.from_bytes(data[pos : pos + 2], "big")
        pos += 2

        # Decompress node names
        nodes = []
        for _ in range(num_nodes):
            if pos + 2 > len(data):
                raise ValueError("SDG.decompress() data truncated")
            name_len = int.from_bytes(data[pos : pos + 2], "big")
            pos += 2
            if pos + name_len > len(data):
                raise ValueError("SDG.decompress() data truncated")
            node = data[pos : pos + name_len].decode("utf-8")
            pos += name_len
            nodes.append(node)

        # Decompress number of edges
        if pos + 2 > len(data):
            raise ValueError("SDG.decompress() data truncated")
        num_edges = int.from_bytes(data[pos : pos + 2], "big")
        pos += 2

        # Build edge type lookup
        edge_type_lookup = {et.value[0]: et for et in EdgeType}

        # Decompress edges
        edges = []
        for _ in range(num_edges):
            if pos + 5 > len(data):
                raise ValueError("SDG.decompress() data truncated")
            source_idx = int.from_bytes(data[pos : pos + 2], "big")
            pos += 2
            target_idx = int.from_bytes(data[pos : pos + 2], "big")
            pos += 2
            edge_type_id = int.from_bytes(data[pos : pos + 1], "big")
            pos += 1

            if source_idx >= len(nodes) or target_idx >= len(nodes):
                raise ValueError("SDG.decompress() invalid node index")
            if edge_type_id not in edge_type_lookup:
                raise ValueError("SDG.decompress() invalid edge type")

            edge_type = edge_type_lookup[edge_type_id]
            edges.append(
                (nodes[source_idx], edge_type.value[3], nodes[target_idx])
            )

        return cls(nodes, edges)
