from typing import Dict, List, Union

import numpy as np
import pandas as pd
from scipy.sparse import coo_matrix
from scipy.spatial import cKDTree

from napari.layers.utils.layer_utils import _FeatureTable
from napari.utils.events.custom_types import Array
from napari.utils.translations import trans


def connex(vertices: np.ndarray) -> list:
    """Connection array to build vertex edges for vispy LineVisual.

    Notes
    -----
    See
    http://api.vispy.org/en/latest/visuals.html#vispy.visuals.LineVisual

    """
    return [True] * (vertices.shape[0] - 1) + [False]


class TrackManager:
    """Manage track data and simplify interactions with the Tracks layer.

    Attributes
    ----------
    data : array (N, D+1)
        Coordinates for N points in D+1 dimensions. ID,T,(Z),Y,X. The first
        axis is the integer ID of the track. D is either 3 or 4 for planar
        or volumetric timeseries respectively.
    features : Dataframe-like
        Features table where each row corresponds to a point and each column
        is a feature.
    properties : dict {str: array (N,)}, DataFrame
        Properties for each point. Each property should be an array of length N,
        where N is the number of points.
    graph : dict {int: list}
        Graph representing associations between tracks. Dictionary defines the
        mapping between a track ID and the parents of the track. This can be
        one (the track has one parent, and the parent has >=1 child) in the
        case of track splitting, or more than one (the track has multiple
        parents, but only one child) in the case of track merging.
        See examples/tracks_3d_with_graph.py
    ndim : int
        Number of spatiotemporal dimensions of the data.
    max_time: float, int
        Maximum value of timestamps in data.
    track_vertices : array (N, D)
        Vertices for N points in D dimensions. T,(Z),Y,X
    track_connex : array (N,)
        Connection array specifying consecutive vertices that are linked to
        form the tracks. Boolean
    track_times : array (N,)
        Timestamp for each vertex in track_vertices.
    graph_vertices : array (N, D)
        Vertices for N points in D dimensions. T,(Z),Y,X
    graph_connex : array (N,)
        Connection array specifying consecutive vertices that are linked to
        form the graph.
    graph_times : array (N,)
        Timestamp for each vertex in graph_vertices.
    track_ids : array (N,)
        Track ID for each vertex in track_vertices.
    """

    def __init__(self) -> None:
        # store the raw data here
        self._data = None
        self._feature_table = _FeatureTable()
        self._order = None

        # use a kdtree to help with fast lookup of the nearest track
        self._kdtree = None

        # NOTE(arl): _tracks and _connex store raw data for vispy
        self._points = None
        self._points_id = None
        self._points_lookup = None
        self._ordered_points_idx = None

        self._track_vertices = None
        self._track_connex = None

        self._graph = None
        self._graph_vertices = None
        self._graph_connex = None

        # lookup table for vertex indices from track id
        self._id2idxs = None

    @staticmethod
    def _fast_points_lookup(sorted_time: np.ndarray) -> Dict[int, slice]:
        """Computes a fast lookup table from time to their respective points slicing."""

        # finds where t transitions to t + 1
        transitions = np.nonzero(sorted_time[:-1] - sorted_time[1:])[0] + 1
        start = np.insert(transitions, 0, 0)

        # compute end of slice
        end = np.roll(start, -1)
        end[-1] = len(sorted_time)

        # access first position of each t slice
        time = sorted_time[start]

        return {t: slice(s, e) for s, e, t in zip(start, end, time)}

    @property
    def data(self) -> np.ndarray:
        """array (N, D+1): Coordinates for N points in D+1 dimensions."""
        return self._data

    @data.setter
    def data(self, data: Union[list, np.ndarray]):
        """set the vertex data and build the vispy arrays for display"""

        # convert data to a numpy array if it is not already one
        data = np.asarray(data)

        # check check the formatting of the incoming track data
        data = self._validate_track_data(data)

        # Sort data by ID then time
        self._order = np.lexsort((data[:, 1], data[:, 0]))
        self._data = data[self._order]

        # build the indices for sorting points by time
        self._ordered_points_idx = np.argsort(self.data[:, 1])
        self._points = self.data[self._ordered_points_idx, 1:]

        # build a tree of the track data to allow fast lookup of nearest track
        self._kdtree = cKDTree(self._points)

        # make the lookup table
        # NOTE(arl): it's important to convert the time index to an integer
        # here to make sure that we align with the napari dims index which
        # will be an integer - however, the time index does not necessarily
        # need to be an int, and the shader will render correctly.
        time = np.round(self._points[:, 0]).astype(np.uint)
        self._points_lookup = self._fast_points_lookup(time)

        # make a second lookup table using a sparse matrix to convert track id
        # to the vertex indices
        self._id2idxs = coo_matrix(
            (
                np.broadcast_to(1, self.track_ids.size),  # just dummy ones
                (self.track_ids, np.arange(self.track_ids.size)),
            )
        ).tocsr()

    @property
    def features(self):
        """Dataframe-like features table.

        It is an implementation detail that this is a `pandas.DataFrame`. In the future,
        we will target the currently-in-development Data API dataframe protocol [1].
        This will enable us to use alternate libraries such as xarray or cuDF for
        additional features without breaking existing usage of this.

        If you need to specifically rely on the pandas API, please coerce this to a
        `pandas.DataFrame` using `features_to_pandas_dataframe`.

        References
        ----------
        .. [1]: https://data-apis.org/dataframe-protocol/latest/API.html
        """
        return self._feature_table.values

    @features.setter
    def features(
        self,
        features: Union[Dict[str, np.ndarray], pd.DataFrame],
    ) -> None:
        self._feature_table.set_values(features, num_data=len(self.data))
        self._feature_table.reorder(self._order)
        if 'track_id' not in self._feature_table.values:
            self._feature_table.values['track_id'] = self.track_ids

    @property
    def properties(self) -> Dict[str, np.ndarray]:
        """dict {str: np.ndarray (N,)}: Properties for each track."""
        return self._feature_table.properties()

    @properties.setter
    def properties(self, properties: Dict[str, Array]):
        """set track properties"""
        self.features = properties

    @property
    def graph(self) -> Dict[int, Union[int, List[int]]]:
        """dict {int: list}: Graph representing associations between tracks."""
        return self._graph

    @graph.setter
    def graph(self, graph: Dict[int, Union[int, List[int]]]):
        """set the track graph"""
        self._graph = self._validate_track_graph(graph)

    @property
    def track_ids(self):
        """return the track identifiers"""
        return self.data[:, 0].astype(np.uint32)

    @property
    def unique_track_ids(self):
        """return the unique track identifiers"""
        return np.unique(self.track_ids)

    def __len__(self):
        """return the number of tracks"""
        return len(self.unique_track_ids) if self.data is not None else 0

    def _vertex_indices_from_id(self, track_id: int):
        """return the vertices corresponding to a track id"""
        return self._id2idxs[track_id].nonzero()[1]

    def _validate_track_data(self, data: np.ndarray) -> np.ndarray:
        """validate the coordinate data"""

        if data.ndim != 2:
            raise ValueError(
                trans._('track vertices should be a NxD array', deferred=True)
            )

        if data.shape[1] < 4 or data.shape[1] > 5:
            raise ValueError(
                trans._(
                    'track vertices should be 4 or 5-dimensional',
                    deferred=True,
                )
            )

        # check that all IDs are integers
        ids = data[:, 0]
        if not np.all(np.floor(ids) == ids):
            raise ValueError(
                trans._('track id must be an integer', deferred=True)
            )

        if not all(t >= 0 for t in data[:, 1]):
            raise ValueError(
                trans._(
                    'track timestamps must be greater than zero', deferred=True
                )
            )

        return data

    def _validate_track_graph(
        self, graph: Dict[int, Union[int, List[int]]]
    ) -> Dict[int, List[int]]:
        """validate the track graph"""

        # check that graph nodes are of correct format
        for node_idx, parents_idx in graph.items():
            # make sure parents are always a list
            if type(parents_idx) != list:
                graph[node_idx] = [parents_idx]

        unique_track_ids = set(self.unique_track_ids)

        # check that graph nodes exist in the track id lookup
        for node_idx, parents_idx in graph.items():
            nodes = [node_idx, *parents_idx]
            for node in nodes:
                if node not in unique_track_ids:
                    raise ValueError(
                        trans._(
                            'graph node {node_idx} not found',
                            deferred=True,
                            node_idx=node_idx,
                        )
                    )

        return graph

    def build_tracks(self):
        """build the tracks"""

        points_id = []
        track_vertices = []
        track_connex = []

        # NOTE(arl): this takes some time when the number of tracks is large
        for idx in self.unique_track_ids:
            indices = self._vertex_indices_from_id(idx)

            # grab the correct vertices and sort by time
            vertices = self.data[indices, 1:]

            # coordinates of the text identifiers, vertices and connections
            points_id += [idx] * vertices.shape[0]
            track_vertices.append(vertices)
            track_connex.append(connex(vertices))

        self._points_id = np.array(points_id)[self._ordered_points_idx]
        self._track_vertices = np.concatenate(track_vertices, axis=0)
        self._track_connex = np.concatenate(track_connex, axis=0)

    def build_graph(self):
        """build the track graph"""

        graph_vertices = []
        graph_connex = []

        for node_idx, parents_idx in self.graph.items():
            # we join from the first observation of the node, to the last
            # observation of the parent
            node_start = self._vertex_indices_from_id(node_idx)[0]
            node = self.data[node_start, 1:]

            for parent_idx in parents_idx:
                parent_stop = self._vertex_indices_from_id(parent_idx)[-1]
                parent = self.data[parent_stop, 1:]

                graph_vertices.append([node, parent])
                graph_connex.append([True, False])

        # if there is a graph, store the vertices and connection arrays,
        # otherwise, clear the vertex arrays
        if graph_vertices:
            self._graph_vertices = np.concatenate(graph_vertices, axis=0)
            self._graph_connex = np.concatenate(graph_connex, axis=0)
        else:
            self._graph_vertices = None
            self._graph_connex = None

    def vertex_properties(self, color_by: str) -> np.ndarray:
        """return the properties of tracks by vertex"""

        if color_by not in self.properties:
            raise ValueError(
                trans._(
                    'Property {color_by} not found',
                    deferred=True,
                    color_by=color_by,
                )
            )

        return self.properties[color_by]

    def get_value(self, coords):
        """use a kd-tree to lookup the ID of the nearest tree"""
        if self._kdtree is None:
            return None

        # query can return indices to points that do not exist, trim that here
        # then prune to only those in the current frame/time
        # NOTE(arl): I don't like this!!!
        d, idx = self._kdtree.query(coords, k=10)
        idx = [i for i in idx if i >= 0 and i < self._points.shape[0]]
        pruned = [i for i in idx if self._points[i, 0] == coords[0]]

        # if we have found a point, return it
        if pruned and self._points_id is not None:
            return self._points_id[pruned[0]]
        return None  # return the track ID

    @property
    def ndim(self) -> int:
        """Determine number of spatiotemporal dimensions of the layer."""
        return self.data.shape[1] - 1

    @property
    def max_time(self) -> int:
        """Determine the maximum timestamp of the dataset"""
        return int(np.max(self.track_times))

    @property
    def track_vertices(self) -> np.ndarray:
        """return the track vertices"""
        return self._track_vertices

    @property
    def track_connex(self) -> np.ndarray:
        """vertex connections for drawing track lines"""
        return self._track_connex

    @property
    def graph_vertices(self) -> np.ndarray:
        """return the graph vertices"""
        return self._graph_vertices

    @property
    def graph_connex(self):
        """vertex connections for drawing the graph"""
        return self._graph_connex

    @property
    def track_times(self) -> np.ndarray:
        """time points associated with each track vertex"""
        return self.track_vertices[:, 0]

    @property
    def graph_times(self) -> np.ndarray:
        """time points associated with each graph vertex"""
        if self.graph_vertices is not None:
            return self.graph_vertices[:, 0]
        return None

    def track_labels(self, current_time: int) -> tuple:
        """return track labels at the current time"""
        # this is the slice into the time ordered points array
        if current_time not in self._points_lookup:
            return [], []

        lookup = self._points_lookup[current_time]
        pos = self._points[lookup, ...]
        lbl = [f'ID:{i}' for i in self._points_id[lookup]]
        return lbl, pos
