import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.cluster import DBSCAN
from sklearn.mixture import GaussianMixture
from sklearn.metrics.pairwise import euclidean_distances


def representative_selection(data, N_represenative = 40, PCA_N_components = 5, DBSCAN_eps = 0.5, DBSCAN_min_samples = 5):
    """
    The representative selection procedure
    it run 1. PCA to original data
           2. DBSCAN for cluster phase 1
           3. Gaussian mixture model on each cluster in phase 1
           4. Find the point with smallest euclidean distance to these means
    and returns the index of these points in original data

    input:
        data: numpy.ndarray, the dataset of image, in current setting, each image need to be flattened (sklearn.PCA only takes flattened value)
        N_represenative: int, the minimal value of representative, the final result may have higher number
        PCA_N_components: int, the number of component for PCA
        DBSCAN_eps: float, the eps hyperparameter fpr DBSCAN
        DBSCAN_min_samples: int

    output:
        rep_list: list of int, the index of representatives in the dataset
                  (which can be use as index of np.array directly)
    """
    #scale all the value to 0-1 scale
    data /= np.max(data)

    #run PCA on X
    pca = PCA(n_components=PCA_N_components)
    X_extracted = pca.fit_transform(data)
    #run DBSCAN on the principle components, get the label
    clustering_phase_1 = DBSCAN(eps = DBSCAN_eps, min_samples = DBSCAN_min_samples)
    clustering_phase_1.fit(X_extracted)
    cluster_phase_1_label = clustering_phase_1.labels_

    #calculate the distribution based on the cluster p1
    rep_distribution = get_rep_distribution(N_represenative = N_represenative, label = cluster_phase_1_label)

    #get the actual sub cluster
    cluster_p1 = get_sub_cluster_p1(data = X_extracted, label = cluster_phase_1_label)

    rep_list = []

    #for each cluster
    for i in range(len(cluster_p1)):

        sub_cluster = cluster_p1[i]
        num_rep = rep_distribution[i]

        #run a GMM to find the mean
        GMM = GaussianMixture(n_components = num_rep)
        GMM.fit(sub_cluster)
        means = GMM.means_

        #find the closest image
        reps = find_closest(cluster = sub_cluster, means = means)
        #add them to the closest image
        rep_list += list(np.where(X_extracted == rep)[0][1] for rep in reps)

    return rep_list

def get_rep_distribution(N_represenative, label):
    """
    given the cluster label generated by DBSCAN, calculate the number of representative for each cluster

    input:
        label: numpy.ndarray, the output of sklearn.cluster.DBSCAN's label_ method
    output:
        rep_distribution: numpy.ndarray, the number of representative for each cluster, in the same order
    """
    label_, count = np.unique(label, return_counts = True)
    #remove the noise term(cluster = -1)
    count = count[1:]
    rep_distribution = np.ceil(count*N_represenative/count.sum()).astype("int")
    return rep_distribution


def get_sub_cluster_p1(data, label):
    """
    given the original dataset and the label from sklearn.cluster, return the list of each cluster

    input:
        data: numpy.ndarray, original data passed into sklearn.cluster.DBSCAN's fit method
        label: numpy.ndarray, the output of sklearn.cluster.DBSCAN's label_ method
    output:
        cluster_p1: list of numpy.ndarray, the feature divided by labels
    """
    unique_label = np.unique(label)[1:]
    cluster_p1 = [[] for i in np.arange(unique_label.shape[0])]
    for i in np.arange(np.shape(label)[0]):
        if (label[i] != -1):
            cluster_p1[label[i]].append(data[i])
    cluster_p1 = [np.array(i) for i in cluster_p1]
    return cluster_p1

def find_closest(cluster, means):
    """
    find the single sample(image) in the cluster which has smallest euclidean distance to mean

    input:
        cluster: numpy.ndarray, a set of sample(image)
        mean: numpy.ndarray, an array of means, from the output of sklearn.mixture.GMM.means_

    output:
        rep: numpy.ndarray, the sample in the cluster which has smallest euclidean distance to each mean respectively
    """

    rep_index = np.argmin(euclidean_distances(cluster, means), axis = 0)
    reps = cluster[rep_index]
    return reps
