import numpy as np
import warnings
from random import randint
from typing import List, Tuple
from sklearn import preprocessing
from tqdm.auto import tqdm
from typing import List, Optional, Union
from doc_utils import DocUtils
from relevanceai.logger import LoguruLogger
from relevanceai.api.client import BatchAPIClient
from relevanceai.vector_tools.cluster import KMeans
from relevanceai.data_tools.base_text_processing import BaseTextProcessing


class PlotTextThemeModel(BatchAPIClient, BaseTextProcessing, LoguruLogger, DocUtils):
    def __init__(
        self,
        project: str,
        api_key: str,
        firebase_uid: str,
        # dataset_info
        dataset_id: str,
        upload_chunksize: int = 50,
        # clustering
        cluster_field: str = "_cluster_",
        # dimensionality reduction parameters
        embedding_dims: int = 2,
        dim_red_k: int = 800,
        n_epochs_without_progress: int = 100,
        # nlp parameters
        language: str = "english",
    ):
        self.dataset_id = dataset_id
        self.upload_chunksize = upload_chunksize
        self.cluster_field = cluster_field
        self.embedding_dims = embedding_dims
        self.dim_red_k = dim_red_k
        self.n_epochs_without_progress = n_epochs_without_progress
        self.language = language
        super().__init__(project=project, api_key=api_key, firebase_uid=firebase_uid)

    def _build_and_plot_clusters(
        self,
        vector_fields: List[str],
        text_fields: List[str],
        max_doc_num: int = None,
        k: int = 10,
        alias: str = "kmeans",
        lower: bool = True,
        remove_digit: bool = True,
        remove_punct: bool = True,
        remove_stop_words: bool = True,
        additional_stop_words: Optional[List[str]] = None,
        cluster_representative_cnt: int = 3,
        plot_axis: str = "off",
        figsize: Tuple[int, ...] = (20, 10),
        cmap: str = "plasma",
        alpha: float = 0.2,
    ):
        additional_stop_words = (
            [] if additional_stop_words is None else additional_stop_words
        )

        # get documents
        documents = self._get_documents(
            vector_fields=vector_fields,
            text_fields=text_fields,
            max_doc_num=max_doc_num,
        )

        # perform kmeans clustering
        alias = alias + "_" + str(k)
        centers, clustered_documents = self._kmeans_clustering(
            documents=documents, vector_fields=vector_fields, k=k, alias=alias
        )

        clustering_results = self.get_field_across_documents(
            ".".join([self.cluster_field, ".".join(vector_fields), alias]), documents
        )
        for i, cl_doc in enumerate(tqdm(clustered_documents)):
            documents[i][
                "_".join([self.cluster_field, "".join(vector_fields), alias])
            ] = clustering_results[i]

        # get cluster data
        cluster_data = self._get_cluster_datafield(
            documents=documents,
            vector_fields=vector_fields,
            text_fields=text_fields,
            alias=alias,
            lower=lower,
            remove_digit=remove_digit,
            remove_punct=remove_punct,
        )
        clusters_top_n_words = self._get_cluster_word_freq(
            cluster_data=cluster_data,
            remove_stop_words=remove_stop_words,
            additional_stop_words=additional_stop_words,
            cluster_representative_cnt=cluster_representative_cnt,
        )
        for c in clusters_top_n_words:
            self.logger.info(c, clusters_top_n_words[c])

        # dimensionality reduction
        # 1) Fit the model and reduce vector size
        vectors: List = [[] for _ in documents]
        for vector_field in vector_fields:
            t = self.get_field_across_documents(vector_field, documents)
            vectors = [i + j for i, j in zip(vectors, t)]  # concatenate vectors
        dr_documents = self._dim_reduction(
            vector_data=vectors,
            embedding_dims=self.embedding_dims,
            k=self.dim_red_k,
            n_epochs_without_progress=self.n_epochs_without_progress,
        )
        for i, dr in enumerate(dr_documents):
            documents[i][vector_field + "_dr_vector_"] = dr.tolist()
        # 2) reduce center vector size
        vectors = [c["centroid_vector_"] for c in centers]
        dr_centers = self._dim_reduction(
            vector_data=vectors,
            embedding_dims=self.embedding_dims,
            k=self.dim_red_k,
            n_epochs_without_progress=self.n_epochs_without_progress,
            run_fit=False,
        )
        for i, c in enumerate(centers):
            c["centroid_dr_vector_"] = dr_centers[i]

        # plot
        self._plot_clusters(
            documents,
            dr_documents,
            centers,
            cluster_data,
            vector_fields,
            alias,
            cluster_representative_cnt,
            plot_axis=plot_axis,
            figsize=figsize,
            cmap=cmap,
            alpha=alpha,
        )

    def _get_documents(
        self, vector_fields: List[str], text_fields: List[str], max_doc_num: int = None
    ):
        self.logger.info(" * Loading documents")
        filters = [
            {
                "field": vector_field,
                "filter_type": "exists",
                "condition": "==",
                "condition_value": " ",
            }
            for vector_field in vector_fields
        ]
        fields = vector_fields + text_fields
        if max_doc_num:
            page_size = 200 if max_doc_num > 200 else max_doc_num
            batch_doc = self._batch_load_documents(
                fields=fields, filters=filters, page_size=page_size
            )
            documents = batch_doc["documents"]
            cursor = batch_doc["cursor"]
            while (
                batch_doc["documents"] != []
                and len(documents) > max_doc_num - page_size
            ):
                batch_doc = self._batch_load_documents(
                    fields=fields, filters=filters, page_size=page_size, cursor=cursor
                )
                documents.extend(batch_doc["documents"])
        else:
            documents = self.get_all_documents(
                dataset_id=self.dataset_id, filters=filters, select_fields=fields
            )
        return documents

    def _batch_load_documents(
        self,
        fields: List[str],
        filters: Optional[List[dict]] = None,
        page_size: int = 200,
        cursor: str = None,
    ):
        return self.datasets.documents.get_where(
            dataset_id=self.dataset_id,
            select_fields=fields,
            page_size=page_size,
            filters=[] if filters is None else filters,
            cursor=cursor,
        )

    def _kmeans_clustering(
        self,
        documents: List[dict],
        vector_fields: List[str],
        k: int = 10,
        alias: str = "kmeans",
    ):
        self.logger.info(" * Kmeans Clustering")
        clusterer = KMeans(k=k)
        clustered_documents = clusterer.fit_documents(
            vector_fields, documents, alias=alias
        )
        res = self._update_documents(
            self.dataset_id, clustered_documents, chunksize=self.upload_chunksize
        )
        centers = clusterer.get_centroid_documents()
        self.services.cluster.centroids.insert(
            dataset_id=self.dataset_id,
            cluster_centers=centers,
            vector_fields=vector_fields,
            alias=alias,
        )
        return centers, clustered_documents

    @staticmethod
    def _get_fields_value_by_id(documents: List[dict], id: str, fields: List[str]):
        vals = {}
        for doc in documents:
            if doc["_id"] == id:
                for field in fields:
                    vals[field] = DocUtils.get_field(field, doc) + " "
        return vals

    def _get_cluster_datafield(
        self,
        documents: List[dict],
        vector_fields: List[str],
        text_fields: List[str],
        alias: str,
        lower: bool = True,
        remove_digit: bool = True,
        remove_punct: bool = True,
    ):
        self.logger.info(" * Extracting cluster info")
        cluster_data: dict = {}
        for i, doc in tqdm(enumerate(documents)):
            cluster_name = doc[
                "_".join([self.cluster_field, "".join(vector_fields), alias])
            ]
            if cluster_name not in cluster_data:
                cluster_data[cluster_name] = {"data": []}
            values = self._get_fields_value_by_id(
                documents, id=doc["_id"], fields=text_fields
            )
            text = self.normalize_text(
                txt=" ".join([v for k, v in values.items()]),
                lower=lower,
                remove_digit=remove_digit,
                remove_punct=remove_punct,
            )
            cluster_data[cluster_name]["data"].append(text)
        return cluster_data

    @staticmethod
    def _get_cluster_population(cluster_data: dict):
        return [
            (cluster_name, len(v["data"])) for cluster_name, v in cluster_data.items()
        ]

    def _get_cluster_word_freq(
        self,
        cluster_data: dict,
        remove_stop_words: bool = True,
        additional_stop_words: List[str] = [],
        cluster_representative_cnt: int = 3,
    ):
        for cluster in tqdm(cluster_data):
            cluster_data[cluster]["word_freq"] = self.get_word_frequency(
                str_list=cluster_data[cluster]["data"],
                remove_stop_words=remove_stop_words,
                additional_stop_words=[]
                if additional_stop_words is None
                else additional_stop_words,
                language=self.language,
            )
        return {
            c: [cluster_data[c]["word_freq"][:cluster_representative_cnt]]
            for c in cluster_data
        }

    def _dim_reduction(
        self,
        vector_data: Union[List, np.ndarray],
        embedding_dims=2,
        k=800,
        n_epochs_without_progress=100,
        run_fit=True,
    ):

        try:
            from ivis import Ivis
        except ModuleNotFoundError:
            warnings.warn(
                """You are missing ivis, please run the appropriate installation option:
            `pip install ivis[gpu]` If you have CUDA installed
            `pip install ivis[cpu]` if you don't have CUDA installed
            """
            )

        self.logger.info(" * Dimensionality reduction")
        vector_data = np.array(vector_data)
        if run_fit:
            self.dr_model = Ivis(
                embedding_dims=embedding_dims,
                k=k,
                n_epochs_without_progress=n_epochs_without_progress,
            )
            dr_documents = self.dr_model.fit_transform(vector_data)
        else:
            dr_documents = self.dr_model.transform(vector_data)

        return dr_documents

    def _plot_clusters(
        self,
        documents: List[dict],
        dr_documents: List,
        centers: List[dict],
        cluster_data: dict,
        vector_fields: List[str],
        alias: str,
        cluster_representative_cnt: int,
        plot_axis: str = "off",
        figsize: Tuple[int, ...] = (20, 10),
        cmap: str = "plasma",
        alpha: float = 0.2,
    ):
        group = [
            x["_".join([self.cluster_field, "".join(vector_fields), alias])]
            for x in documents
        ]
        le = preprocessing.LabelEncoder()
        color = le.fit_transform(group)

        try:
            import matplotlib.pyplot as plt
        except ModuleNotFoundError:
            warnings.warn(
                "You are missing matplotlib, please run `pip install matplotlib`"
            )

        plt.axis(plot_axis)
        plt.figure(figsize=figsize)
        plt.scatter(
            [x[0] for x in dr_documents],
            [x[1] for x in dr_documents],
            c=color,
            cmap=cmap,
            alpha=alpha,
        )

        for i, r in enumerate(centers):
            plt.annotate(
                [
                    x[0]
                    for x in cluster_data[r["_id"]]["word_freq"][
                        :cluster_representative_cnt
                    ]
                ],
                (r["centroid_dr_vector_"][0] - 0.5, r["centroid_dr_vector_"][1]),
            )


def build_and_plot_clusters(
    self,
    project: str,
    api_key: str,
    firebase_uid: str,
    dataset_id: str,
    vector_fields: List[str],
    text_fields: List[str],
    upload_chunksize: int = 50,
    cluster_field: str = "_cluster_",
    embedding_dims: int = 2,
    dim_red_k: int = 800,
    n_epochs_without_progress: int = 100,
    language: str = "english",
    max_doc_num: int = None,
    k: int = 10,
    alias: str = "kmeans",
    lower: bool = True,
    remove_digit: bool = True,
    remove_punct: bool = True,
    remove_stop_words: bool = True,
    additional_stop_words: Optional[List[str]] = None,
    cluster_representative_cnt: int = 3,
    plot_axis: str = "off",
    figsize: Tuple[int, ...] = (20, 10),
    cmap: str = "plasma",
    alpha: float = 0.2,
):
    model = PlotTextThemeModel(
        project=project,
        api_key=api_key,
        firebase_uid=firebase_uid,
        dataset_id=dataset_id,
        upload_chunksize=upload_chunksize,
        # clustering
        cluster_field=cluster_field,
        # dimensionality reduction parameters
        embedding_dims=embedding_dims,
        dim_red_k=dim_red_k,
        n_epochs_without_progress=n_epochs_without_progress,
        language=language,
    )
    return model._build_and_plot_clusters(
        vector_fields=vector_fields,
        text_fields=text_fields,
        max_doc_num=max_doc_num,
        k=k,
        alias=alias,
        lower=lower,
        remove_digit=remove_digit,
        remove_punct=remove_punct,
        remove_stop_words=remove_stop_words,
        additional_stop_words=[]
        if additional_stop_words is None
        else additional_stop_words,
        cluster_representative_cnt=cluster_representative_cnt,
        plot_axis=plot_axis,
        figsize=figsize,
        cmap=cmap,
        alpha=alpha,
    )
