from collections.abc import Iterable, Mapping, Sequence
from typing import Any, Literal, TypeVar, overload

import networkx as nx
import numpy as np
from array_api._2024_12 import Array
from array_api_compat import array_namespace
from jacobi_poly import lgamma
from strenum import StrEnum

TCartesian = TypeVar("TCartesian")
TSpherical = TypeVar("TSpherical")
COSSIN = Literal["cos", "sin"]


class BranchingType(StrEnum):
    """
    The branching types of the nodes in a rooted tree representing the coordinates.

    (Vilenkin's method of trees)
    """

    A = "a"
    B = "b"
    BP = "b'"
    C = "c"


def check_tree(graph: nx.DiGraph, /) -> None:
    """
    Check if the graph is a valid tree representing the coordinates.
    Definition of "a tree representing the coordinates":
    - The graph is a tree.
    - The graph is rooted.
    - Each node has 0 or 2 successors.
      (full binary tree, not necessarily perfect binary tree)
    - If the node has 2 successors, one of the outgoing edges has
      attribute "type" of "cos" and
      the other has attribute "type" of "sin".

    Parameters
    ----------
    graph : nx.DiGraph
        The graph representing the coordinates.

    Raises
    ------
    ValueError
        If the graph is not a valid tree.

    """
    # check if the graph is a tree
    if not nx.is_tree(graph):
        raise ValueError("The graph is not a tree.")
    # check if the graph is rooted
    if not nx.is_arborescence(graph):
        raise ValueError("The graph is not rooted.")
    # check every node has 0 or 2 successors
    for node in graph.nodes:
        if (degree := graph.out_degree(node)) not in {0, 2}:
            raise ValueError(
                f"Each node should have 0 or 2 successors, "
                f"but {node} has {degree} successors."
            )
    # check every node except the root has attribute
    # "direction" of type Literal["cos", "sin"]
    for node in graph.nodes:
        if graph.out_degree(node) == 0:
            continue
        if {e[2] for e in graph.out_edges(nbunch=node, data="type")} != {
            "cos",
            "sin",
        }:
            raise ValueError(
                f"Node {node} should have 0 or 2 successors "
                "with types 'cos' and 'sin', "
                f"got {set(graph.out_edges.data(nbunch=node, data='type'))}."
            )


def get_non_leaf_descendants(graph: nx.DiGraph, /) -> dict[Any, int]:
    """
    Calculate the number of non-leaf descendants for each node in a rooted tree.

    Parameters
    ----------
    graph : nx.DiGraph
        A rooted tree.

    Returns
    -------
    dict[Any, int]
        NaN for leaf nodes, otherwise the number of non-leaf descendants.

    """
    # calculate number of non-leaf descendants for each node
    leaf_nodes = {node for node in graph.nodes if graph.out_degree(node) == 0}
    non_leaf_descendants = {}
    for node in reversed(list(nx.topological_sort(graph))):
        if node in leaf_nodes:
            non_leaf_descendants[node] = -1
        else:
            non_leaf_descendants[node] = sum(
                non_leaf_descendants[successor] + 1
                for successor in graph.successors(node)
            )
    for node in leaf_nodes:
        non_leaf_descendants[node] = np.nan
    return non_leaf_descendants


def get_child(G: nx.DiGraph, node: Any, type: COSSIN, /) -> Any:
    """
    Get the child node of the given type in a rooted tree representing the coordinates.

    Parameters
    ----------
    G : nx.DiGraph
        A rooted tree representing the coordinates.
    node : Any
        The node to get the child.
    type : COSSIN
        The type of the child._description_

    Raises
    ------
    ValueError
        If the node has no child of the given type.
        If the node has multiple children of the given type.

    """
    child_candidates = [
        child for child in G.successors(node) if G.edges[node, child]["type"] == type
    ]
    if len(child_candidates) == 0:
        raise ValueError(f"No {type} child for node {node}.")
    elif len(child_candidates) > 1:
        raise ValueError(f"Multiple {type} children for node {node}.")
    return child_candidates[0]


def get_parent(G: nx.DiGraph, node: Any, /) -> Any | None:
    """
    Get the parent node in a rooted tree representing the coordinates.

    Parameters
    ----------
    G : nx.DiGraph
        A rooted tree representing the coordinates.
    node : Any
        The node to get the parent.

    Returns
    -------
    Any
        The parent node if exists, otherwise None.

    Raises
    ------
    ValueError
        If the node has multiple parents.

    """
    predecessors = list(G.predecessors(node))
    if len(predecessors) == 0:
        return None
    elif len(predecessors) > 1:
        raise ValueError("The node has multiple parents.")
    return predecessors[0]


def get_branching_types(G: nx.DiGraph, /) -> dict[Any, BranchingType]:
    """
    Get the branching types of each node in a rooted tree.

    Parameters
    ----------
    G : nx.DiGraph
        A rooted tree.

    Returns
    -------
    dict[Any, BranchingType]
        The branching types of each node.

    """
    branching_types = {}
    for node in G.nodes:
        if G.out_degree(node) == 0:
            continue
        sin_child = get_child(G, node, "sin")
        cos_child = get_child(G, node, "cos")
        sin_child_is_leaf = G.out_degree(sin_child) == 0
        cos_child_is_leaf = G.out_degree(cos_child) == 0
        if cos_child_is_leaf and sin_child_is_leaf:
            branching_types[node] = BranchingType.A
        elif cos_child_is_leaf and not sin_child_is_leaf:
            branching_types[node] = BranchingType.B
        elif not cos_child_is_leaf and sin_child_is_leaf:
            branching_types[node] = BranchingType.BP
        else:
            branching_types[node] = BranchingType.C
    return branching_types


def get_branching_type_from_digraph(G: nx.DiGraph, /) -> Iterable[BranchingType]:
    """
    Get the branching types of each node in a rooted tree.

    Parameters
    ----------
    G : nx.DiGraph
        A rooted tree.

    Returns
    -------
    Sequence[BranchingType]
        The branching types of each node.

    """
    root = next(nx.topological_sort(G))
    stack = [root]
    while stack:
        node = stack.pop()
        if G.out_degree(node) == 0:
            raise AssertionError()
        cos_child = get_child(G, node, "cos")
        sin_child = get_child(G, node, "sin")
        cos_child_is_leaf = G.out_degree(cos_child) == 0
        sin_child_is_leaf = G.out_degree(sin_child) == 0
        if cos_child_is_leaf and sin_child_is_leaf:
            branching_type = BranchingType.A
        elif cos_child_is_leaf and not sin_child_is_leaf:
            branching_type = BranchingType.B
            stack.append(sin_child)
        elif not cos_child_is_leaf and sin_child_is_leaf:
            branching_type = BranchingType.BP
            stack.append(cos_child)
        else:
            branching_type = BranchingType.C
            stack.extend([sin_child, cos_child])
        yield branching_type


class SphericalCoordinates[TSpherical, TCartesian]:
    """Stores the spherical coordinates using the method of trees by Vilenkin.

    Examples
    --------
    >>> import ultrasphere as us
    >>> c = us.create_standard(3)
    >>> c
    SphericalCoordinates(bba)
    >>> c.s_nodes
    ['theta0', 'theta1', 'theta2']
    >>> c.c_nodes
    [0, 1, 2, 3]
    >>> c.branching_types
    {'theta0': <BranchingType.B: 'b'>, 'theta1': <BranchingType.B: 'b'>, 'theta2': <BranchingType.A: 'a'>}
    >>> c.cos_edges
    [('theta0', 0), ('theta1', 1), ('theta2', 2)]
    >>> c.sin_edges
    [('theta0', 'theta1'), ('theta1', 'theta2'), ('theta2', 3)]
    >>> c.S
    {3: nan, 2: nan, 'theta2': 0, 1: nan, 'theta1': 1, 0: nan, 'theta0': 2}
    """

    G: nx.DiGraph
    """The rooted tree representing the coordinates."""
    S: Mapping[TSpherical, int]
    """The number of non-leaf descendants for each spherical node."""
    branching_types: Mapping[TSpherical, BranchingType]
    """The branching types of each node."""
    s_nodes: Sequence[TSpherical]
    """The spherical nodes."""
    c_nodes: Sequence[TCartesian]
    """The Cartesian nodes."""
    cos_edges: Sequence[tuple[TSpherical, TCartesian | TSpherical]]
    """The edges with type 'cos'."""
    sin_edges: Sequence[tuple[TSpherical, TCartesian | TSpherical]]
    """The edges with type 'sin'."""

    def __hash__(self) -> int:
        """__hash__."""
        return hash(self.G)

    def __getstate__(self) -> dict[Any, Any]:
        """__getstate__."""
        return {"G": nx.to_dict_of_dicts(self.G)}

    def __setstate__(self, state: dict[Any, Any]) -> None:
        """__setstate__."""
        return self.__init__(nx.from_dict_of_dicts(state["G"]))  # type: ignore

    def __eq__(self, other: Any) -> bool:
        """Check if the spherical coordinates are equal."""
        if not isinstance(other, SphericalCoordinates):
            return False
        return self.G == other.G

    def __repr__(self) -> str:
        """Print the spherical coordinates."""
        return f"SphericalCoordinates({self.branching_types_expression_str})"

    def __str__(self) -> str:
        """Print the spherical coordinates."""
        return f"SphericalCoordinates({self.branching_types_expression_str})"

    @property
    def root(self) -> Any:
        """The root node.

        Examples
        --------
        >>> import ultrasphere as us
        >>> us.create_standard(3).root
        'theta0'
        """
        return next(nx.topological_sort(self.G))

    @property
    def root_index(self) -> int:
        """The index of the root node."""
        return self.s_nodes.index(self.root)

    @property
    def branching_types_expression(self) -> Sequence[BranchingType]:
        """The branching types.

        Examples
        --------
        >>> import ultrasphere as us
        >>> us.create_standard(2).branching_types_expression
        [<BranchingType.B: 'b'>, <BranchingType.A: 'a'>]
        """
        return list(get_branching_type_from_digraph(self.G))

    @property
    def branching_types_expression_str(self) -> str:
        """The branching types as a string.

        Examples
        --------
        >>> import ultrasphere as us
        >>> us.create_standard(3).branching_types_expression_str
        'bba'
        """
        return "".join(
            branching_type.value for branching_type in self.branching_types_expression
        )

    @property
    def s_ndim(self) -> int:
        """The number of spherical dimensions.

        Examples
        --------
        >>> import ultrasphere as us
        >>> us.create_standard(3).s_ndim
        3"""
        return len(self.s_nodes)

    @property
    def c_ndim(self) -> int:
        """The number of Cartesian dimensions.

        Examples
        --------
        >>> import ultrasphere as us
        >>> us.create_standard(3).c_ndim
        4"""
        return len(self.c_nodes)

    def __init__(self, tree: nx.DiGraph, /) -> None:
        """
        Initialize the spherical coordinates.

        Parameters
        ----------
        tree : nx.DiGraph
            The rooted tree representing the coordinates.

            Definition of "a tree representing the coordinates":
            - The graph is a tree.
            - The graph is rooted.
            - Each node has 0 or 2 successors.
            (full binary tree, not necessarily perfect binary tree)
            - If the node has 2 successors, one of the outgoing edges has
            attribute "type" of "cos" and
            the other has attribute "type" of "sin".

        """
        check_tree(tree)
        self.G = tree
        self.S = get_non_leaf_descendants(tree)
        self.branching_types = get_branching_types(tree)
        nodes = list(nx.topological_sort(tree))
        # even if sorted by name, c_nodes are still topologically sorted
        # because they do not have any successors
        self.c_nodes = sorted(node for node in nodes if tree.out_degree(node) == 0)
        self.s_nodes = [node for node in nodes if tree.out_degree(node) == 2]
        self.cos_edges = [e for e in self.G.edges if self.G.edges[e]["type"] == "cos"]
        self.sin_edges = [e for e in self.G.edges if self.G.edges[e]["type"] == "sin"]

    def surface_area(self, r: float = 1) -> float:
        r"""
        The surface area of the unit sphere.

        .. math::
            |\mathbb{S}^{d-1}| = \frac{2 \pi^{d/2}}{\Gamma(d/2)} r^{d-1}

        Returns
        -------
        float
            The surface area.

        Examples
        --------
        >>> import ultrasphere as us
        >>> from array_api_compat import numpy as np
        >>> np.round(us.create_standard(2).surface_area(), 6)
        np.float64(12.566371)

        """
        return (
            2
            * (np.pi ** (self.c_ndim / 2))
            / np.exp(lgamma(self.c_ndim / 2.0))
            * r**self.s_ndim
        )

    def volume(self, r: float = 1) -> float:
        r"""
        The volume of the unit sphere.

        .. math::
            \Upsilon_d = \frac{\pi^{d/2}}{\Gamma(d/2 + 1)} r^d

        Returns
        -------
        float
            The volume.

        References
        ----------
        McLean, W. (2000). Strongly Elliptic Systems and
        Boundary Integral Equations. p.247 (Upsilon_n)

        Examples
        --------
        >>> import ultrasphere as us
        >>> from array_api_compat import numpy as np
        >>> np.round(us.create_standard(2).volume(), 6)
        np.float64(4.18879)
        """
        return (
            (np.pi ** (self.c_ndim / 2))
            / np.exp(lgamma(self.c_ndim / 2.0 + 1.0))
            * r**self.c_ndim
        )

    def from_cartesian(
        self, cartesian: Mapping[TCartesian, Array], /
    ) -> Mapping[TSpherical | Literal["r"], Array]:
        """
        Convert the Cartesian coordinates to the spherical coordinates.

        Parameters
        ----------
        cartesian : Mapping[TCartesian, Array]
            The Cartesian coordinates.

        Returns
        -------
        Array
            The spherical coordinates.

        Examples
        --------
        >>> import ultrasphere as us
        >>> from array_api_compat import numpy as np
        >>> c = us.create_spherical()
        >>> s = c.from_cartesian(np.asarray([1.0, 2.0, 3.0]))
        >>> {k: np.round(s[k], 6) for k in c.s_nodes + ["r"]}
        {'theta': np.float64(0.640522), 'phi': np.float64(1.107149), 'r': np.float64(3.741657)}

        """
        xp = array_namespace(*[cartesian[k] for k in self.c_nodes])
        r = (
            xp.linalg.vector_norm(
                xp.stack(xp.broadcast_arrays(*[cartesian[k] for k in self.c_nodes])),
                axis=0,
            )
            if self.s_ndim > 0
            else cartesian[self.c_nodes[0]]
        )
        result: dict[TSpherical | Literal["r"], Array] = {"r": r}
        tmp = {k: cartesian[k] for k in self.c_nodes}
        for node in reversed(list(nx.topological_sort(self.G))):
            if node in self.c_nodes:
                continue
            elif node not in self.s_nodes:
                raise AssertionError()
            cos_child = get_child(self.G, node, "cos")
            sin_child = get_child(self.G, node, "sin")
            cos = tmp[cos_child]
            sin = tmp[sin_child]
            result[node] = xp.atan2(sin, cos)
            tmp[node] = xp.linalg.vector_norm(
                xp.stack(xp.broadcast_arrays(sin, cos)), axis=0
            )
        return result

    @overload
    def to_cartesian(
        self,
        spherical: (
            Mapping[TSpherical | Literal["r"], Array] | Mapping[TSpherical, Array]
        ),
        /,
        *,
        as_array: Literal[False] = ...,
    ) -> Mapping[TCartesian, Array]: ...

    @overload
    def to_cartesian(
        self,
        spherical: (
            Mapping[TSpherical | Literal["r"], Array] | Mapping[TSpherical, Array]
        ),
        /,
        *,
        as_array: Literal[True] = ...,
    ) -> Array: ...

    def to_cartesian(
        self,
        spherical: (
            Mapping[TSpherical | Literal["r"], Array] | Mapping[TSpherical, Array]
        ),
        /,
        *,
        as_array: bool = False,
    ) -> Mapping[TCartesian, Array] | Array:
        """
        Convert the spherical coordinates to the Cartesian coordinates.

        Parameters
        ----------
        spherical : Mapping[TSpherical, Array]
            The spherical coordinates.
        as_array : bool, optional
            Whether to return as an array, by default False

        Returns
        -------
        Array
            The Cartesian coordinates.

        Examples
        --------
        >>> import ultrasphere as us
        >>> from array_api_compat import numpy as np
        >>> c = us.create_spherical()
        >>> e = c.to_cartesian({"theta": np.float64(0.640522), "phi": np.float64(1.107149), "r": np.float64(3.741657)})
        >>> {k: np.round(e[k], 5) for k in c.c_nodes}
        {0: np.float64(1.0), 1: np.float64(2.0), 2: np.float64(3.0)}
        """
        xp = array_namespace(*[spherical[k] for k in self.s_nodes])
        result = {self.root: spherical.get("r", 1)}  # type: ignore
        for node in nx.topological_sort(self.G):
            if node in self.c_nodes:
                continue
            elif node not in self.s_nodes:
                raise AssertionError()
            cos_child = get_child(self.G, node, "cos")
            sin_child = get_child(self.G, node, "sin")
            cos = xp.cos(spherical[node])
            sin = xp.sin(spherical[node])
            result[cos_child] = result[node] * cos
            result[sin_child] = result[node] * sin
        result = {k: result[k] for k in self.c_nodes}
        if as_array:
            return xp.stack(
                xp.broadcast_arrays(*[result[k] for k in self.c_nodes]),
                axis=0,
            )

        return result

    @property
    def is_c_keys_range(self) -> bool:
        """Whether the Cartesian keys are 0, 1, ..., c_ndim - 1."""
        return set(self.c_nodes) == set(range(self.c_ndim))

    @property
    def is_s_keys_range(self) -> bool:
        """Whether the spherical keys are 0, 1, ..., s_ndim - 1."""
        return set(self.s_nodes) == set(range(self.s_ndim))

    def jacobian(
        self,
        spherical: (
            Mapping[TSpherical | Literal["r"], Array] | Mapping[TSpherical, Array]
        ),
        /,
    ) -> Array:
        """
        Calculate the Jacobian of the spherical coordinates.

        Parameters
        ----------
        spherical : Mapping[TSpherical, Array]
            The spherical coordinates.

        Returns
        -------
        Array
            The Jacobian of the spherical coordinates.

        """
        xp = array_namespace(*[spherical[k] for k in self.s_nodes])
        if "r" in spherical:
            jacobian = spherical["r"]  # type: ignore
        else:
            jacobian = 1
        for node in self.s_nodes:
            cos = xp.cos(spherical[node])
            sin = xp.sin(spherical[node])
            cos_child = get_child(self.G, node, "cos")
            sin_child = get_child(self.G, node, "sin")
            jacobian *= cos ** (self.S.get(cos_child, -1) + 1) * sin ** (
                self.S.get(sin_child, -1) + 1
            )
        return jacobian
