import gpbasics.global_parameters as global_param

global_param.ensure_init()

from sklearn.cluster import DBSCAN
from typing import List
import numpy as np
import gpbasics.KernelBasics.PartitioningModel as pm
import gpbasics.DataHandling.DataInput as di
import gpbasics.Metrics.Metrics as met
import time
import logging


class ClusterCenterCriterion(pm.PartitionCriterion):
    def __init__(self, cluster_center: np.ndarray):
        assert len(cluster_center.shape) == 1, "Invalid cluster center given: shape=%s. " \
                                               "Shape of cluster center has to be [d,]." % str(cluster_center.shape)
        super(ClusterCenterCriterion, self).__init__(pm.PartitioningClass.SMALLEST_DISTANCE)
        self.cluster_center: np.ndarray = cluster_center

    def get_score(self, x_vector: np.ndarray) -> np.ndarray:
        return np.sqrt(np.sum(np.square(x_vector - self.cluster_center), axis=1))

    def deepcopy(self):
        return ClusterCenterCriterion(self.cluster_center.copy())

    def get_json(self) -> dict:
        return {"type": "cluster_center", "center": self.cluster_center.tolist()}


class DbscanModel(pm.PartitioningModel):
    def __init__(self, ignored_dimensions: List[int]):
        super(DbscanModel, self).__init__(pm.PartitioningClass.SMALLEST_DISTANCE, ignored_dimensions)
        self.max_window_size = 500

    def automatic_init_criteria(self, data_input: di.DataInput, optimize_metric: met.MetricType,
                                model_selection_metric: met.MetricType):
        logging.info("Setting up Automatically determining Partitioning Criteria by KMeans")
        clusters_k = int(data_input.n_train / self.max_window_size)
        filtered_x_train = self.filter_data_by_ignored_dimensions(data_input.data_x_train)

        logging.info("Starting determining Partitioning Criteria by KMeans. k=%i, shape=%s"
                     % (clusters_k, str(filtered_x_train.shape)))

        eps_factor = float(np.abs(clusters_k * 2))
        min_samples = float(np.abs((data_input.n_train / self.max_window_size)))
        start_time = time.time()
        eps = (np.sum(np.max(data_input.data_x_train) - np.min(data_input.data_x_train))) / eps_factor
        dbscan = DBSCAN(eps=eps, min_samples=min_samples).fit(filtered_x_train)

        logging.info("Finished partitioning in %f s" % (time.time() - start_time))

        partition_sizes = []

        partition_sizes.append(len(dbscan.labels_[dbscan.labels_ == -1]))
        for i in range(np.max(dbscan.labels_) + 1):
            partition_sizes.append(len(dbscan.labels_[dbscan.labels_ == i]))

        logging.info("Partition (#%i) sizes statistics: min %f, max %f, mean %f, median %f"
                     % (len(partition_sizes), np.min(partition_sizes), np.max(partition_sizes),
                        float(np.mean(partition_sizes)), float(np.median(partition_sizes))))

        partitioning: List[ClusterCenterCriterion] = []

        logging.info("Initializing partition criteria by means of previously determined cluster centers.")
        for cluster_center in dbscan.cluster_centers_:
            partitioning.append(ClusterCenterCriterion(cluster_center))

        self.init_partitioning(partitioning)

    def deepcopy(self):
        partitioning: List[ClusterCenterCriterion] = [pc.deepcopy() for pc in self.partitioning]
        copied_self = DbscanModel(self.ignored_dimensions.copy())
        copied_self.init_partitioning(partitioning)
        return copied_self
