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

class Clusterer:
    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:
        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()

    def extract_location_features(self, locations: List[dict]) -> np.ndarray:
        return np.asarray([(float(l['latitude']), float(l['longitude'])) for l in locations])

    def extract_time_features(self, times: List[Dict]) -> np.ndarray:
        return np.asarray([((t['timestamp']), 0) for t in times])

    def label_dataset(self, dataset:List[Dict], labels:List) -> List:
        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)):
            dataset[i]['cluster_label'] = labels[i]

    def group_by_clusters(self, dataset:List[Dict], labels:List) -> Dict[int, List[Dict]]:
        clusters = {}
        for label in labels:
            clusters[label] = [ds for ds in dataset if ds['cluster_label'] == label]
        
        return clusters

    def cluster_locations(self, locations:List[Dict]) -> Dict[int, List[Dict]]:
        '''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)
        
    def cluster_times(self, times:List[Dict]) -> Dict[int, List[Dict]]:
        '''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)