"""

A module to implement useful function to handle trees.

Trees are stored as dictionaries.

"""

import json

import warnings

from typing import List, Optional, Sequence, Tuple, Union

import numpy as np

from nptyping import NDArray, Number, Shape

from scipy.sparse import spmatrix

from sklearn.neighbors import NearestNeighbors

from .errors import (

    InvalidEdgeProps,

    InvalidNodeProps,

    InvalidTree,

    UnknownNodeProperty,

    UnrelatedNode,

)

from .kruskal import UFDS

from .types import (

    BinaryNodeProperty,

    Childhood,

    Edge,

    EdgeProperties,

    Node,

    NodeProperties,

    NumericalEdgeProperty,

    NumericalNodeProperty,

    Parenthood,

)

def complete_tree(

    parents: Optional[Parenthood] = None, children: Optional[Childhood] = None

):

    if parents is None:

        parents = {}

        if children is None:

            children = {}

        else:

            for parent in children:

                children[parent] = set(children[parent])

                for child in children[parent]:

                    parents[child] = parent

    else:

        if children is None:

            children = {}

            for child, parent in parents.items():

                try:

                    children[parent].add(child)

                except KeyError:

                    children[parent] = {child}

        else:

            for parent in children:

                children[parent] = set(children[parent])

                for child in children[parent]:

                    if child not in parents or parents[child] != parent:

                        raise InvalidTree

    return parents, children

def get_root(parents: Parenthood, node: Node = None) -> Node:

    """

    Get root of a node in a tree.

    *****************************

    """

    if node is None:

        for k, v in parents.items():

            node = k

            return get_root(parents, k)

    if node not in parents:

        return node

    root = node

    while root in parents:

        root = parents[root]

    return root

def get_root_path(parents: Parenthood, node: Node) -> List[Node]:

    """

    Get path to root of a node in a tree.

    *************************************

    """

    if node not in parents:

        warnings.warn("Requested node {} is not in the parenthood.".format(node))

        return [node]

    root = node

    root_path = [node]

    while root in parents:

        root = parents[root]

        root_path.append(root)

    return root_path

def get_root_path_match(parents: Parenthood, node: Node, target: Node) -> List[Node]:

    """

    Get path to root of a node in a tree.

    *************************************

    """

    if target not in parents:

        warnings.warn("Target node {} is not in the parenthood.".format(target))

        return get_root_path(parents, node)

    if node not in parents:

        warnings.warn("Requested node {} is not in the parenthood.".format(node))

        return []

    root = node

    root_path = [node]

    while root in parents:

        if root == target:

            return root_path

        root = parents[root]

        root_path.append(root)

def get_leaves_without_prop(

    children: Childhood,

    node: Node,

) -> List[Node]:

    """

    Get leaves of a node in a tree.

    *******************************

    """

    if node not in children:

        return [node]

    no_lvs = [node]

    lvs = []

    while len(no_lvs) > 0:

        new_no_lvs = []

        for n in no_lvs:

            if n in children:

                for c in children[n]:

                    new_no_lvs.append(c)

            else:

                lvs.append(n)

        no_lvs = new_no_lvs

    return lvs

def get_leaves_with_prop(

    children: Childhood, node: Node, prop: BinaryNodeProperty

) -> List[Node]:

    """

    Get leaves of a node in a tree.

    *******************************

    """

    if node not in prop:

        warnings.warn("Node {} does not have the property".format(node))

        return []

    if not prop[node]:

        warnings.warn("Root {} does not pass the property test.".format(node))

        return []

    if node not in children:

        warnings.warn("Children of Root {} does not pass the property.".format(node))

        return [node]

    no_lvs = [node]

    lvs = []

    while len(no_lvs) > 0:

        new_no_lvs = []

        for n in no_lvs:

            if n in prop:

                if prop[n]:

                    if n in children:

                        candidates = []

                        for c in children[n]:

                            if c in prop:

                                if prop[c]:

                                    candidates.append(c)

                        new_no_lvs += candidates

                        if len(candidates) == 0:

                            lvs.append(n)

                    else:

                        lvs.append(n)

            else:

                warnings.warn("Node {} does not have the property".format(node))

        no_lvs = new_no_lvs

    return lvs

def get_leaves(

    children: Childhood, node: Node, prop: Optional[BinaryNodeProperty] = None

) -> List[Node]:

    """

    Get leaves of a node in a tree.

    *******************************

    """

    if prop is None:

        return get_leaves_without_prop(children, node)

    return get_leaves_with_prop(children, node, prop)

def kruskal_edges(

    edges: Sequence[Edge], weights: NumericalEdgeProperty

) -> Sequence[Edge]:

    """

    Yield kruskal edges, given a list of edges.

    ********************************************

    """

    # create Union-Find data structure

    components = UFDS()

    # edges are sorted by non-decreasing order of dissimilarity

    edges = sorted(edges, key=lambda x: weights[x])

    k_edges = []

    k_weights = []

    for edge in edges:

        # nodes in involved in the edge

        n1, n2 = edge

        # roots of nodes in the Union-Find

        rn1 = components.get_root(n1)

        rn2 = components.get_root(n2)

        # if components are differents

        if rn1 != rn2:

            components.union(edge)

            k_edges.append(edge)

            k_weights.append(weights[edge])

    return k_edges, k_weights

def kruskal_tree(

    edges: Sequence[Edge], weights: NumericalEdgeProperty, size: NumericalNodeProperty

) -> Tuple[Parenthood, Childhood, NumericalNodeProperty]:

    """

    Create parents an children relationships from kruskal edges.

    ***********************************************************

    """

    parents = dict()

    children = dict()

    props = {"weights": dict(), "size": dict()}

    k_edges, k_weights = kruskal_edges(edges, weights)

    max_node = 2 * len(k_edges)

    for edge, weight in zip(k_edges, k_weights):

        n1, n2 = edge

        rn1 = get_root(parents, n1)

        rn2 = get_root(parents, n2)

        if rn1 in props["size"]:

            s1 = props["size"][rn1]

        else:

            s1 = size[n1]

            props["size"][rn1] = s1

        if rn2 in props["size"]:

            s2 = props["size"][rn2]

        else:

            s2 = size[n2]

            props["size"][rn2] = s2

        # since it is already a spanning tree,

        # I know rn1 and rn2 have different roots

        parents[rn1] = max_node

        parents[rn2] = max_node

        children[max_node] = [rn1, rn2]

        props["weights"][max_node] = weight

        props["size"][max_node] = s1 + s2

        max_node += 1

    return parents, children, props

def tree_to_json(

    nodes: Sequence[Node],

    parents: Parenthood,

    children: Childhood,

    jsonfile: str,

    nodeprops: Optional[NodeProperties] = None,

    edgeprops: Optional[EdgeProperties] = None,

):

    """Store a jsonified tree to a json file."""

    output_dict = dict()

    output_dict["nodes"] = nodes

    output_dict["parents"] = parents

    output_dict["children"] = children

    output_dict["nodeprops"] = dict()

    output_dict["edgeprops"] = dict()

    if nodeprops is not None:

        if isinstance(nodeprops, dict):

            for k, v in nodeprops.items():

                output_dict["nodeprops"][k] = v

        else:

            raise InvalidNodeProps(

                "Invalid node props, "

                "expected {} but got {}".format(dict, type(nodeprops))

            )

    if edgeprops is not None:

        if isinstance(edgeprops, dict):

            for k, v in edgeprops.items():

                output_dict["edgeprops"][k] = v

        else:

            raise InvalidEdgeProps(

                "Invalid node props, "

                "expected {} but got {}".format(dict, type(edgeprops))

            )

    json_dict = json.dumps(output_dict)

    with open(jsonfile, "w") as outputjson:

        outputjson.write(json_dict)

def _expand_on_property(

    cut: List[Node],

    children: Childhood,

    prop: NumericalNodeProperty,

    threshold: Union[int, float],

) -> List[Node]:

    """Create a new tree by cutting based on property threshold."""

    candidates = []

    expansion = []

    for node in cut:

        if node in children:

            candidates += children[node]

    for candidate in candidates:

        if candidate in prop:

            if prop[candidate] >= threshold:

                expansion.append(candidate)

    return expansion

def cut_on_property(

    parents: Parenthood,

    children: Childhood,

    prop: NumericalNodeProperty,

    threshold: Union[int, float],

) -> List[Node]:

    """Produce a list of authorized nodes given a property threshold."""

    root = get_root(parents)

    cut = set()

    remaining = [root]

    while len(remaining) > 0:

        cut |= set(remaining)

        remaining = _expand_on_property(remaining, children, prop, threshold)

    return list(cut)

def common_ancestor(parents: Parenthood, node1: Node, node2: Node) -> Node:

    """Get the common ancestor of two nodes and store their distances to him."""

    if node1 in parents and node2 in parents:

        rp1 = get_root_path(parents, node1)

        rp2 = get_root_path(parents, node2)

        if len(set(rp1) & set(rp2)) > 0:

            for node in rp1:

                if node in rp2:

                    return node

        raise UnrelatedNode(

            "Nodes {} and {} have no common ancestors!!!".format(node1, node2)

        )

    raise UnrelatedNode(

        "One of the provided nodes: ({}, {}) has no parent...".format(node1, node2)

    )

def edge_dist(parents: Parenthood, node1: Node, node2: Node) -> int:

    """Return the number of edges to go from node1 to node2 (by common ancestor)."""

    ancestor = common_ancestor(parents, node1, node2)

    rpm1 = get_root_path_match(parents, node1, ancestor)

    rpm2 = get_root_path_match(parents, node2, ancestor)

    return len(set(rpm1) | set(rpm2))

def weighted_dist(

    parents: Parenthood, weights: NumericalNodeProperty, node1: Node, node2: Node

) -> float:

    """Return the number of edges to go from node1 to node2 (by common ancestor)."""

    ancestor = common_ancestor(parents, node1, node2)

    rpm1 = get_root_path_match(parents, node1, ancestor)

    rpm2 = get_root_path_match(parents, node2, ancestor)

    nodes_in_path = set(rpm1) | set(rpm2)

    nodes_in_path.discard(node1)

    nodes_in_path.discard(node2)

    dist = 0.0

    for node in nodes_in_path:

        if node not in weights:

            raise UnknownNodeProperty(

                "Missing weight for node {} to compute a weighted distance!!!".format(

                    node

                )

            )

        dist += weights[node]

    # minus 1 otherwise ancestor is counted twice

    return dist

def farthest_point_sampling(

    coords: NDArray[Shape["N_points, N_dims"], Number], n_samples: Union[int, float]

) -> NDArray[Shape["N_samples"], np.int32]:

    """

    Perform farthest points sampling using point coordinates array.

    Args:

        coords: array containing point coordinates.

        n_samples: number of point to sample. If a float is given, represents the

            proportion of points used instead.

    Returns:

        Array containing idxs of sampled points.

    """

    if isinstance(n_samples, float):

        n_samples = int(n_samples * len(coords))

    idxs = np.zeros(n_samples, dtype=np.int32)

    idxs[0] = np.random.randint(len(coords))

    distances = ((coords[idxs[0]] - coords) ** 2).sum(1)

    for i in range(1, n_samples):

        idxs[i] = np.argmax(distances)

        distances = np.minimum(distances, ((coords[idxs[i]] - coords) ** 2).sum(1))

    return idxs

def random_farthest_point_sampling(

    coords: NDArray[Shape["N_points, N_dims"], Number],

    n_farthest_samples: Union[int, float] = 0.3,

    n_random_samples: Union[int, float] = 0.1,

) -> NDArray[Shape["N_samples"], np.int32]:

    """

    Perform farthest points sampling using point coordinates array followed by random

    sampling .

    Args:

        coords: array containing point coordinates.

        n_farthest_samples: number of points to keep using farthest points sampling.

            If a float is given, represents the proportion of points used instead.

        n_random_samples: number of points to keep using random sampling. If a float

            is given, represents the proportion of points used instead.

    Returns:

        Array containing idxs of sampled points.

    """

    farthest_idxs = farthest_point_sampling(coords, n_farthest_samples)

    if isinstance(n_random_samples, float):

        n_random_samples = int(n_random_samples * len(coords))

    probs = np.ones(len(coords))

    probs[farthest_idxs] = 0

    random_idxs = np.arange(len(coords))

    random_idxs = np.random.choice(random_idxs, size=n_random_samples, p=probs)

    idxs = np.concatenate((farthest_idxs, random_idxs))

    return idxs

def get_kneighbors_graph(

    points: NDArray[Shape["N_points, N_dims"], Number],

    n_farthest_samples: Union[int, float] = 0.3,

    n_random_samples: Union[int, float] = 0.1,

    dmax: int = 500,

    n_neighbors: int = 5,

    n_jobs: Optional[int] = None,

) -> spmatrix:

    """

    Get a graph generated by KNN on given points.

    Args:

        points: array containg point coordinates.

        n_farthest_samples: number of points to keep using farthest points sampling. If

            a float is given, represents the proportion of points used instead.

        n_random_samples: number of points to keep using random sampling. If a float is

            given, represents the proportion of points used instead.

        dmax: maximum distance in pixels between two adjacent nodes.

        n_neighbors: number of neighbors to use for KNN algorithm.

        n_jobs: number of parallel jobs to run for neighbors search. None means 1.

    Returns:

        Sparse distance matrix representing the graph.

    """

    idxs = random_farthest_point_sampling(

        points,

        n_farthest_samples=n_farthest_samples,

        n_random_samples=n_random_samples,

    )

    X = points[idxs]

    knn = NearestNeighbors(n_neighbors=n_neighbors, n_jobs=n_jobs).fit(X)

    A = knn.kneighbors_graph(mode="distance")

    Abool = A.astype(bool) - (A > dmax)

    A = A.multiply(Abool)

    return A.maximum(A.T)

def get_nodeprops_edgeprops(

    A: spmatrix, coords: NDArray[Shape["N_points, N_dims"], Number]

) -> Tuple[NodeProperties, EdgeProperties]:

    """

    Get coordinates and distances between edges of a graph as NodeProperties and

    EdgeProperties.

    Args:

        A: Sparse distance matrix representing the graph.

        coords: coordinates of the nodes.

    Returns:

        NodeProperties dictionary containing 'x' and 'y' entries for node coordinates

        and EdgeProperties dictionary containing a 'distance' entry for distances

        between edges.

    """

    edgeprops = {"distance": {(i, j): A[i, j] for i, j in zip(*A.nonzero())}}

    nodeprops = {"x": {}, "y": {}}

    for i, (x, y) in enumerate(coords):

        nodeprops["x"][i] = x

        nodeprops["y"][i] = y

    return nodeprops, edgeprops