from typing import List
from .config import BaselineConfig, Configuration
from ..utils import __create_model__
# import numpy as np
from sklearn.cluster import KMeans
# from yellowbrick.cluster import KElbowVisualizer
from .clusters import ClusterList
class ClusterPipeline:
    def __init__(self, config:Configuration = None):
        if config is None:
            self.__setup__(BaselineConfig())
        else:
            self.__setup__(config)

    def __setup__(self, config:Configuration):
        self.PTM = __create_model__(config.plm)
        self.dimension_reduction = __create_model__(config.dimension_reduction) # TODO
        self.clustering = __create_model__(config.clustering)
        self.keywords_extraction = __create_model__(config.keywords_extraction)

    def __1_generate_word_embeddings__(self, documents: List[str]):
        '''

        :param documents: a list of N strings:
        :return: np.ndarray: Nx384 (sentence-transformers)
        '''
        print(f'>>> start generating word embeddings...')
        print(f'>>> successfully generated word embeddings...')
        return self.PTM.encode(documents)

    def __2_dimenstion_reduction__(self, embeddings):
        '''

        :param embeddings: NxD
        :return: Nxd, d<<D
        '''
        if self.dimension_reduction is None:
            return embeddings
        print(f'>>> start dimension reduction...')
        print(f'>>> finished dimension reduction...')

    def __3_clustering__(self, embeddings, return_cluster_centers = False, best_k: int = 5):
        '''

        :param embeddings: Nxd
        :return:
        '''
        if self.clustering is None:
            return embeddings
        else:
            print(f'>>> start clustering...')
            model = KMeans()
            # visualizer = KElbowVisualizer(
            #     model, k=(2, 12), metric='calinski_harabasz', timings=False, locate_elbow=False
            # )
            #
            # visualizer.fit(embeddings)
            # best_k = visualizer.k_values_[np.argmax(np.array(visualizer.k_scores_))]
            # print(f'>>> The best K is {best_k}.')

            labels, cluster_centers = self.clustering(embeddings, k=best_k)
            clusters = ClusterList(best_k)
            clusters.instantiate(labels)
            print(f'>>> finished clustering...')

            if return_cluster_centers:
                return clusters, cluster_centers
            return clusters

    def __4_keywords_extraction__(self, clusters: ClusterList, documents: List[str]):
        '''

        :param clusters: N documents
        :return: clusters, where each cluster has added keyphrases
        '''
        if self.keywords_extraction is None:
            return clusters
        else:
            print(f'>>> start keywords extraction')
            for cluster in clusters:
                doc_ids = cluster.elements()
                input_abstracts = [documents[i] for i in doc_ids] #[str]
                keyphrases = self.keywords_extraction(input_abstracts) #[{keys...}]
                cluster.add_keyphrase(keyphrases)
                # for doc_id in doc_ids:
                #     keyphrases = self.keywords_extraction(documents[doc_id])
                #     cluster.add_keyphrase(keyphrases)
            print(f'>>> finished keywords extraction')
            return clusters


    def __call__(self, documents: List[str], best_k:int = 5):
        print(f'>>> pipeline starts...')
        x = self.__1_generate_word_embeddings__(documents)
        x = self.__2_dimenstion_reduction__(x)
        clusters = self.__3_clustering__(x,best_k=best_k)
        outputs = self.__4_keywords_extraction__(clusters, documents)
        print(f'>>> pipeline finished!\n')
        return outputs