# import sys

import sys

sys.path.append('..')

# progress bar import

from tqdm import tqdm

# numpy, sklearn imports

import numpy as np

from sklearn.cluster import KMeans

from sklearn.decomposition import PCA

from sklearn.metrics import silhouette_samples, silhouette_score

from tensorflow.keras.preprocessing.image import load_img

# utils imports

from models import *

from utils.dataset import ImageCLEFDataset

class Cluster:

    def __init__(self, K:int, clef_dataset:ImageCLEFDataset):

        """ Class to perform K-Means clustering in ImageCLEF dataset. We used this system for the contest

        Args:

            K (int): The K clusters we want

            clef_dataset (ImageCLEFDataset): The dataset we employed. Only CLEF is acceptable.

        """

        self.K = K

        self.dataset = clef_dataset

    def do_PCA(self, features:dict) -> np.array:

        """ Perforrms Principal Component Analysis (PCA), to reduce the huge size of the arrays

        Args:

            features (dict): The image_ids, image_vectors pairs.

        Returns:

            np.array: The image_ids, image_vectors pairs, with reduced size.

        """

        feat = np.array(list(features.values()))

        feat = feat.reshape(-1, feat.shape[2])

        pca = PCA(n_components=100, random_state=22)

        pca.fit(feat)

        x = pca.transform(feat)

        return x

    def do_Kmeans(self, x:np.array) -> KMeans:

        """ Fit the K-Means

        Args:

            x (np.array): The image vectors

        Returns:

            KMeans: The fitted K-Means object

        """

        kmeans = KMeans(n_clusters=self.K, random_state=22)

        kmeans.fit(x)

        return kmeans

    def load_features(self) ->tuple[list[dict], list[dict], list[dict]]:

        """ Loads train, validation, test sets 

        Returns:

            tuple[list[dict], list[dict], list[dict]]: The train, validation, test sets in dictionary format

        """

        return self.dataset.get_splits_sets()

    def clustering(self) -> tuple[dict, dict, dict]:

        """ Performs the k-Means clustering using the fitted K-Means object.

        Returns:

            tuple[dict, dict, dict]: The clustered train, val, test image_ids, image_vectors pairs.

        """

        # load splits

        train_features, valid_features, test_features = self.load_features()

        # get the ids for each split

        train_ids, val_ids, test_ids = list(train_features[0].keys()), list(valid_features[0].keys()), list(test_features[0].keys())

        # concate all features to perform a more efficient K-Means

        all_features = dict(train_features, **valid_features)

        all_features = dict(all_features, **test_features)

        # reduce size for fast training

        pca = self.do_PCA(all_features)

        # perform clustering

        kmeans = self.do_Kmeans(pca)

        train_index_limit, val_index_limit = len(train_features), len(train_features)+len(valid_features)

        # get the clustering labels for each set

        train_k_means_labels = kmeans.labels_[:train_index_limit]

        valid_k_means_labels = kmeans.labels_[train_index_limit:val_index_limit]

        test_k_means_labels = kmeans.labels_[val_index_limit:]

        print('# train kmeans:',  len(train_k_means_labels))

        print('# dev kmeans:',  len(valid_k_means_labels))

        print('# test kmeans:',  len(test_k_means_labels))

        # store the clustered train, validation, test set images

        groups_train = {}

        for file, cluster in tqdm(zip(train_ids, train_k_means_labels)):

            if cluster not in groups_train.keys():

                groups_train[cluster] = []

                groups_train[cluster].append(file)

            else:

                groups_train[cluster].append(file)

        groups_valid = {}

        for file, cluster in tqdm(zip(val_ids, valid_k_means_labels)):

            if cluster not in groups_valid.keys():

                groups_valid[cluster] = []

                groups_valid[cluster].append(file)

            else:

                groups_valid[cluster].append(file)

        groups_test = {}

        for file, cluster in tqdm(zip(test_ids, test_k_means_labels)):

            if cluster not in groups_test.keys():

                groups_test[cluster] = []

                groups_test[cluster].append(file)

            else:

                groups_test[cluster].append(file)

        print('# train kmeans:',  len(groups_train))

        print('# dev kmeans:',  len(groups_valid))

        print('# test kmeans:',  len(groups_test))

        return groups_train, groups_valid, groups_test