import warnings
from typing import Optional, Union

import numpy as np
from scipy.sparse import dok_array
from sklearn.base import ClusterMixin, TransformerMixin
from sklearn.cluster import HDBSCAN
from sklearn.feature_extraction.text import CountVectorizer

from turftopic import ClusteringTopicModel
from turftopic.base import Encoder
from turftopic.encoders.utils import encode_chunks
from turftopic.vectorizers.phrases import PhraseVectorizer


class C_Top2Vec(ClusteringTopicModel):
    """Convenience function to construct a C-Top2Vec model in Turftopic.
    The model is essentially the same as ClusteringTopicModel
    with defaults that resemble C-Top2Vec (UMAP -> HDBSCAN -> Centroid term importance).

    ```bash
    pip install turftopic[umap-learn]
    ```

    ```python
    from turftopic import C_Top2Vec

    corpus: list[str] = ["some text", "more text", ...]

    model = C_Top2Vec().fit(corpus)
    model.print_topics()
    ```
    """

    def __init__(
        self,
        encoder: Union[
            Encoder,
            str,
        ] = "sentence-transformers/all-MiniLM-L6-v2",
        window_size: int = 50,
        step_size: int = 40,
        batch_size: int = 64,
        vectorizer: Optional[CountVectorizer] = None,
        dimensionality_reduction: Optional[TransformerMixin] = None,
        clustering: Optional[ClusterMixin] = None,
        feature_importance="centroid",
        n_reduce_to: Optional[int] = None,
        reduction_method="smallest",
        reduction_distance_metric="cosine",
        reduction_topic_representation="centroid",
        random_state: Optional[int] = None,
    ):
        self.batch_size = batch_size
        self.window_size = window_size
        self.step_size = step_size
        if dimensionality_reduction is None:
            try:
                from umap import UMAP
            except ModuleNotFoundError as e:
                raise ModuleNotFoundError(
                    "UMAP is not installed in your environment, but Top2Vec requires it."
                ) from e
            dimensionality_reduction = UMAP(
                n_neighbors=50,
                n_components=5,
                min_dist=0.0,
                metric="cosine",
                random_state=random_state,
            )
        if clustering is None:
            clustering = HDBSCAN(
                min_cluster_size=15,
                metric="euclidean",
                cluster_selection_method="eom",
            )
        if vectorizer is None:
            vectorizer = PhraseVectorizer(
                max_ngram=3,
                min_df=10,
                max_df=1.0,
                threshold=0.5,
                stop_words="english",
                smoothing=5,
            )
        super().__init__(
            encoder=encoder,
            vectorizer=vectorizer,
            dimensionality_reduction=dimensionality_reduction,
            clustering=clustering,
            n_reduce_to=n_reduce_to,
            random_state=random_state,
            feature_importance=feature_importance,
            reduction_method=reduction_method,
            reduction_distance_metric=reduction_distance_metric,
            reduction_topic_representation=reduction_topic_representation,
        )

    def produce_chunk_embeddings(self, raw_documents):
        chunk_embeddings, chunk_positions = encode_chunks(
            self.encoder_,
            raw_documents,
            batch_size=self.batch_size,
            window_size=self.window_size,
            step_size=self.step_size,
        )
        return chunk_embeddings, chunk_positions

    def fit_transform(
        self,
        raw_documents,
        y=None,
        embeddings: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        """Fits model and predicts cluster labels for all given documents.

        Parameters
        ----------
        raw_documents: iterable of str
            Documents to fit the model on.
        y: None
            Ignored, when the dimensionality reduction is TSNE (the default),
            in case of a dimensionality reduction that can utilize labels,
            you can pass labels to the model to inform the clustering process.
        embeddings: ndarray of shape (n_documents, n_dimensions), optional
            IGNORED. Chunk embedding has to be done by the model.

        Returns
        -------
        ndarray of shape (n_documents, n_topics)
            Topic proportions in each document.
        """
        if embeddings is not None:
            warnings.warn(
                "Precomputed embeddings are ignored when using C-Top2Vec."
            )
        # Removing unnecessary whitespace because tokenizers don't handle it too well
        raw_documents = [" ".join(doc.split()) for doc in raw_documents]
        chunk_embeddings, chunk_positions = self.produce_chunk_embeddings(
            raw_documents
        )
        doc_labels = []
        flat_embeddings = []
        chunk_text = []
        for i_doc, (emb, pos, doc) in enumerate(
            zip(chunk_embeddings, chunk_positions, raw_documents)
        ):
            n_chunks = len(emb)
            flat_embeddings.extend(emb)
            for start_char, end_char in pos:
                chunk_text.append(doc[start_char:end_char])
            doc_labels.extend([i_doc] * n_chunks)
        doc_labels = np.array(doc_labels)
        flat_embeddings = np.stack(flat_embeddings)
        topic_labels = super().fit_predict(
            chunk_text, embeddings=flat_embeddings
        )
        n_topics = len(self.classes_)
        n_documents = len(raw_documents)
        print("Creating output array")
        doc_topic_matrix = dok_array((n_documents, n_topics), dtype=np.float32)
        label_to_idx = {label: idx for idx, label in enumerate(self.classes_)}
        chunks_per_doc = [len(emb) for emb in chunk_embeddings]
        print("Adding topic membership")
        for i_doc, topic_label, n_chunks in zip(
            doc_labels, topic_labels, chunks_per_doc
        ):
            doc_topic_matrix[i_doc, label_to_idx[topic_label]] += 1 / n_chunks
        return doc_topic_matrix

    def fit_predict(
        self, raw_documents, y=None, embeddings: Optional[np.ndarray] = None
    ) -> np.ndarray:
        """Fits model and predicts cluster labels for all given documents.

        Parameters
        ----------
        raw_documents: iterable of str
            Documents to fit the model on.
        y: None
            Ignored, when the dimensionality reduction is TSNE (the default),
            in case of a dimensionality reduction that can utilize labels,
            you can pass labels to the model to inform the clustering process.
        embeddings: ndarray of shape (n_documents, n_dimensions), optional
            Precomputed document encodings.

        Returns
        -------
        ndarray of shape (n_documents)
            Cluster label for all documents (-1 for outliers)
        """
        return self.classes_[
            np.argmax(
                self.fit_transform(raw_documents, embeddings=embeddings),
                axis=1,
            )
        ]
