import json
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import DBSCAN
from typing import List, Dict, Any, TypeVar
from deprecated import deprecated

T = TypeVar('T')
ClusterGroup = Dict[Any, List[Dict]]

class Clusterer:
    '''
    Clusterer for applying density-based clustering on datasets.
    The clustering is done with DBSCAN.

    :param epsilon: Epsilon used in DBSCAN 
    :param min_points: Min_points used in DBSCAN
    '''
    def __init__(self, epsilon=11, min_points=2):
        self.epsilon = epsilon
        self.min_points = min_points

    def draw_locations(self, locations:List, labels:List=None) -> plt.Figure:
        if locations is None or len(locations) == 0:
            return self._draw_locations()

        if labels is None or len(locations) != len(labels):
            labels = self.create_labels(locations)

        return self._draw_locations(
            locations = self.extract_location_data(locations),
            partition_info = labels
        )

    def _draw_locations(self, locations:np.ndarray=None, centroids:np.ndarray=None, partition_info:List=None) -> plt.Figure:
        fig = plt.Figure()
        axis = fig.add_subplot(1, 1, 1)

        if locations is not None:
            colors = plt.cm.rainbow(np.linspace(0, 1, len(locations)))

            if partition_info is not None:
                distinct_colors = plt.cm.rainbow(np.linspace(0, 1, len(set(partition_info))))
                colors = [distinct_colors[pi] for pi in partition_info]

            # draw locations with random colors
            axis.scatter(locations[:,0],
                        locations[:,1], 
                        c=colors)

        if centroids is not None:
            # draw black centroids
            axis.scatter(centroids[:,0], centroids[:,1], c='k', marker='x', s=80)
            
        return fig
    
    def create_labels(self, features:np.ndarray) -> List[int]:
        '''Creates labels for the items based on DBSCAN.'''
        if features is None or len(features) == 0:
            return features # trash in trash out
    
        dbsc = DBSCAN(eps = self.epsilon, min_samples = self.min_points)
        dbsc = dbsc.fit(features)
        labels = dbsc.labels_

        return labels.tolist()

    @deprecated(reason="Use generic version instead")
    def extract_location_features(self, locations: List[dict]) -> np.ndarray:
        return np.asarray([(float(l['latitude']), float(l['longitude'])) for l in locations])

    @deprecated(reason="Use generic version instead")
    def extract_time_features(self, times: List[Dict]) -> np.ndarray:
        return np.asarray([[float(t['timestamp'])] for t in times])

    def _extract_features(self, dataset: List[Dict], features:List[str]) -> np.ndarray:
        '''Extracts the feature values from the dataset into a np array with same order as original dataset.'''
        extracted_features = []
        for data in dataset:
            entry = [float(data[feature]) for feature in features]
            extracted_features.append(entry)

        return np.asarray(extracted_features)

    def label_dataset(self, dataset:List[Dict], labels:List[Any]) -> List:
        '''Adds the labels to the elements of the dataset at the same position. The new key is called cluster_label.'''
        if dataset is None or labels is None:
            return

        if len(dataset) != len(labels):
            raise ValueError("dataset and labels has to have same length")

        for i in range(len(dataset)):
            if 'cluster_label' in dataset[i]:
                continue
            dataset[i]['cluster_label'] = labels[i]

    def group_by_clusters(self, dataset:List[Dict], labels:List[Any]) -> ClusterGroup:
        self.label_dataset(dataset, labels)

        clusters = {}
        for label in labels:
            clusters[label] = [ds for ds in dataset if ds['cluster_label'] == label]
        
        return clusters

    @deprecated(reason="Use generic version instead")
    def cluster_locations(self, locations:List[Dict]) -> ClusterGroup:
        '''Returns a dictionary with identified clusters and their locations copied from the input'''
        if locations is None or len(locations) == 0:
            # raise Exception("locations has to contain something")
            return {}

        features = self.extract_location_features(locations)
        
        labels = self.create_labels(features)
        self.label_dataset(locations, labels)

        return self.group_by_clusters(locations, labels)
       
    @deprecated(reason="Use generic version instead") 
    def cluster_times(self, times:List[Dict]) -> ClusterGroup:
        '''Returns a dictionary with identified clusters and their times copied from the input'''
        features = self.extract_time_features(times)
    
        labels = self.create_labels(features)
        self.label_dataset(times, labels)

        return self.group_by_clusters(times, labels)

    def cluster_dataset(self, dataset:List[Dict], features:List[str]) -> ClusterGroup:
        '''
        Returns the identified clusters containing a subset of nodes from the dataset.
        
        :param dataset: The nodes to assign to clusters
        :param features: The feature names of the nodes to use for clustering

        :returns: A dictionary of clusters, where each value is a non-empty subset of dataset if dataset was not empty
        '''
        arr = self._extract_features(dataset, features)

        labels = self.create_labels(arr)

        return self.group_by_clusters(dataset, labels)

