from typing import Any, Dict, Optional, Sequence, Tuple, Union

import numpy as np
from nptyping import NDArray, Shape
from scipy.sparse import triu
from sortedcontainers import SortedDict
from tqdm import tqdm

from .object_api import Tree, UGraph
from .types import Edge, Node


class AgglomerativeClustering:
    r"""
    Object used to hierarchically cluster nodes on a graph. Clustering greedily chooses
    to merge linked nodes that have minimum distance/strength ratio. Strength between
    2 nodes is initially 1 for every edge and 0 when there is no edge, then when 2 nodes
    are merged the strength of a newly formed link between the new node and another node
    is the weighted (by node population) average of the strengths between the 2 old
    nodes and the other node. This algorithm uses centroid linkage clustering (UPGMC).

    Args:
        compute_all: whether to initially compute all distances between nodes regardless
            of there linkage.
    """

    def __init__(self, compute_all: bool = False):
        self.compute_all = compute_all

    def init_graph(
        self,
        G: UGraph,
        feats: Union[Dict[Node, NDArray[Shape["*"], Any]], Sequence[str]],
        weights: Optional[Union[Dict[Edge, float], str]] = None,
    ):
        r"""
        Initialize main graph attributes (adjacency matrix, n_nodes and features) using
        a graph object, a list of features and a list of weights.

        Args:
            G: graph to cluster nodes on.
            feats: either a dictionary that maps nodes to their corresponding feature
                vectors or a sequence of property names that will be used as features.
            weights: either a dictionary that maps edges to their corresponding weight
                or a property name that will be used as weight. If `None` is passed,
                weights are computed using euclidian distances between feature vectors.
        """
        self.A = triu(G.A, format="csr").astype(np.float32)
        self.n_nodes = G.n_nodes
        if isinstance(feats, dict):
            feats = [feats[node] for node in G.nodes]
            self.feats = np.stack(feats)
        else:
            self.feats = []
            for node in G.nodes:
                self.feats.append([G.nodeprops[feat][node] for feat in feats])
            self.feats = np.array(self.feats)

        if weights is None:
            ii, jj = self.A.nonzero()
            dists = ((feats[ii] - feats[jj]) ** 2).sum(1)
            self.A[ii, jj] = dists
        elif isinstance(weights, dict):
            for (n1, n2) in weights:
                i, j = sorted((G.nodes.index(n1), G.nodes.index(n2)))
                self.A[i, j] = weights[n1, n2] ** 2
        else:
            for (n1, n2) in G.edges:
                i, j = sorted((G.nodes.index(n1), G.nodes.index(n2)))
                self.A[i, j] = G.edgeprops[str][(n1, n2)] ** 2

    def reset(self):
        """
        Reset the algorithm attributes. Populations are initiated to 1 for every node,
        strengths are initiated to 1 for every edge, dendrogram is emptied.
        """
        self.populations_ = {k: 1 for k in range(self.n_nodes)}
        ii, jj = self.A.nonzero()
        self.centroids_ = {k: self.feats[k] for k in range(self.n_nodes)}
        self.links_ = {
            k: set(jj[ii == k].tolist() + ii[jj == k].tolist())
            for k in range(self.n_nodes)
        }
        self.strengths_ = {(i, j): 1 for i, j in zip(ii, jj)}
        if self.compute_all:
            self.distances_ = {
                (i, j): self.distance(i, j)
                for i in range(self.n_nodes)
                for j in range(i + 1, self.n_nodes)
            }
        else:
            self.distances_ = {(i, j): self.distance(i, j) for i, j in zip(ii, jj)}
        self.edges_ = SortedDict(
            self.criterion, {(i, j): self.distances_[i, j] for i, j in zip(ii, jj)}
        )
        self.dendrogram_ = np.zeros((self.n_nodes - 1, 4))

    def distance(self, i: int, j: int) -> float:
        """
        Get squared distance between nodes `i` and `j`. If available in the adjacency
        matrix or in the `distance` dictionary it is not recomputed.

        Args:
            i: first node.
            j: second node.

        Returns:
            Squared euclidian distance between i and j.
        """
        i, j = sorted((i, j))
        try:
            d = self.A[i, j]
        except IndexError:
            d = self.distances_.get((i, j), 0)
        if not d:
            d = ((self.centroids_[i] - self.centroids_[j]) ** 2).sum()
        return d

    def criterion(self, x: Tuple[int, int]) -> float:
        """
        Criterion function used to find the next nodes to merge. Override it to use
        another criterion.

        Args:
            x: tuple containg the two nodes to merge.

        Returns:
            Squared distance between the 2 nodes divided by link strength.
        """
        i, j = x
        return self.distances_[i, j] / self.strengths_[i, j]

    def create_centroid_link(self, i, j, c, k):
        """
        Create a new link between centroid `c` (that comes from merging nodes `i` and
        `j`) and node `k`.

        Args:
            i: first merged node.
            j: second merged node.
            c: centroid of nodes `i` and `j`.
            k: node linked to either `i` or `j` or both.
        """
        if i == k or j == k:
            return
        ik, ki = sorted((i, k))
        jk, kj = sorted((j, k))
        ck, kc = sorted((c, k))
        ij, ji = sorted((i, j))

        pi = self.populations_[i]
        pj = self.populations_[j]
        ri = pi / (pi + pj)
        rj = pj / (pi + pj)
        try:
            dik = self.distances_[ik, ki]
            djk = self.distances_[jk, kj]
            dij = self.distances_[ij, ji]
            dck = ri * dik + rj * djk - ri * rj * dij
        except KeyError:
            dck = self.distance(c, k)

        self.distances_[ck, kc] = dck

        self.edges_.pop((ik, ki), 0)
        self.edges_.pop((jk, kj), 0)
        sik = self.strengths_.get((ik, ki), 0)
        sjk = self.strengths_.get((jk, kj), 0)
        if sik or sjk:
            self.strengths_[ck, kc] = ri * sik + rj * sjk
            self.edges_[ck, kc] = dck

        self.links_[k].discard(i)
        self.links_[k].discard(j)
        self.links_[k].add(c)
        self.links_[j].discard(k)
        self.links_[c].add(k)

    def add_link(self, i: int, j: int):
        """
        Create a new link between 2 nodes.

        Args:
            i: first node.
            j: second node.
        """
        i, j = sorted((i, j))
        dij = self.distance(i, j)
        self.distances_[i, j] = dij
        self.strengths_[i, j] = 1
        self.edges_[i, j] = dij
        self.links_[i].add(j)
        self.links_[j].add(i)

    def fit(
        self,
        G: UGraph,
        feats: Union[Dict[Node, NDArray[Shape["*"], Any]], Sequence[str]],
        weights: Optional[Union[Dict[Edge, float], str]] = None,
    ):
        r"""
        Fits on the given graph and completes the dendrogram. A dendrogram is an array
        of size :math:`(n-1) \times 4` (whre :math:`n` is the number of nodes)
        representing the successive merges of nodes. Each row gives the two merged
        nodes, their distance and the size of the resulting cluster. Any new node
        resulting from a merge takes the first available index (e.g., the first merge
        corresponds to node :math:`n`).

        Args:
            G: graph to cluster nodes on.
            feats: either a dictionary that maps nodes to their corresponding feature
                vectors or a sequence of property names that will be used as features.
            weights: either a dictionary that maps edges to their corresponding weight
                or a property name that will be used as weight. If `None` is passed,
                weights are computed using euclidian distances between feature vectors.
        """
        self.init_graph(G, feats, weights)
        self.reset()

        c = self.n_nodes
        for n in tqdm(range(self.n_nodes - 1), total=self.n_nodes - 1):
            if not self.edges_:
                cur_dendrogram = self.dendrogram_[:n]
                missing = sorted(
                    [
                        k
                        for k in range(c)
                        if k not in cur_dendrogram[:, 0]
                        and k not in cur_dendrogram[:, 1]
                    ]
                )
                for k, i in enumerate(missing):
                    for j in missing[k + 1 :]:
                        self.add_link(i, j)
            (i, j), _ = self.edges_.popitem(0)

            pi = self.populations_[i]
            pj = self.populations_[j]
            ri = pi / (pi + pj)
            rj = pj / (pi + pj)
            self.dendrogram_[n] = [i, j, self.criterion((i, j)), pi + pj]

            self.centroids_[c] = ri * self.centroids_[i] + rj * self.centroids_[j]
            self.populations_[c] = pi + pj
            self.links_[c] = set()

            while self.links_[i]:
                k = self.links_[i].pop()
                self.create_centroid_link(i, j, c, k)
            while self.links_[j]:
                k = self.links_[j].pop()
                self.create_centroid_link(j, i, c, k)
            self.links_.pop(i)
            self.links_.pop(j)
            c += 1

    def fit_transform(
        self,
        G: UGraph,
        feats: Union[Dict[Node, NDArray[Shape["*"], Any]], Sequence[str]],
        weights: Optional[Union[Dict[Edge, float], str]] = None,
    ) -> Tree:
        """
        Fits on the given graph and returns the hierarchical clustering tree.

        Args:
            G: graph to cluster nodes on.
            feats: either a dictionary that maps nodes to their corresponding feature
                vectors or a sequence of property names that will be used as features.
            weights: either a dictionary that maps edges to their corresponding weight
                or a property name that will be used as weight. If `None` is passed,
                weights are computed using euclidian distances between feature vectors.

        Returns:
            The tree that describes the hierarchical clustering procedure.
        """
        self.fit(G, feats, weights)

        children = []
        parents = []
        nodes = list(range(G.n_nodes))

        if isinstance(weights, str):
            key = weights
        else:
            key = "weight"
        edgeprops = {key: {}}

        if isinstance(feats, dict):
            n_feats = next(iter(feats)).shape
            nodeprops = {
                k: {n: feats[node][k] for n, node in enumerate(G.nodes)}
                for k in range(n_feats)
            }
        else:
            nodeprops = {
                feat: {n: G.nodeprops[feat][node] for n, node in enumerate(G.nodes)}
                for feat in feats
            }
        nodeprops["population"] = {n: 1 for n in range(G.n_nodes)}

        for k, row in enumerate(self.dendrogram_):
            n = k + self.n_nodes
            n1, n2 = row[:2]
            children[n] = [n1, n2]
            parents[n1] = n
            parents[n2] = n
            nodes.append(n)
            if isinstance(feats, dict):
                for k, centroid in enumerate(self.centroids_[n]):
                    nodeprops[k][n] = centroid
            else:
                for k, feat in enumerate(feats):
                    nodeprops[feat][n] = self.centroids_[n, k]
            edgeprops[key][n, n1] = self.distance(n, n1) ** 0.5
            edgeprops[key][n, n2] = self.distance(n, n2) ** 0.5
            nodeprops["population"][n] = self.populations_[n]

        return Tree(nodes, parents, children, nodeprops, edgeprops)