import logging
import time
from typing import Callable, Optional, Tuple, Union

import pandas as pd
import streamlit as st
from bokeh.palettes import Turbo256
from bokeh.plotting import Figure
from embedding_lenses.embedding import embed_text
from embedding_lenses.utils import encode_labels
from embedding_lenses.visualization import draw_interactive_scatter_plot
from sentence_transformers import SentenceTransformer

from perplexity_lenses import REGISTRY_DATASET

logger = logging.getLogger(__name__)
EMBEDDING_MODELS = [
    "distiluse-base-multilingual-cased-v1",
    "distiluse-base-multilingual-cased-v2",
    "all-mpnet-base-v2",
    "flax-sentence-embeddings/all_datasets_v3_mpnet-base",
]
DIMENSIONALITY_REDUCTION_ALGORITHMS = ["UMAP", "t-SNE"]
DOCUMENT_TYPES = ["Whole document", "Sentence"]
SEED = 0
LANGUAGES = [
    "af",
    "ar",
    "az",
    "be",
    "bg",
    "bn",
    "ca",
    "cs",
    "da",
    "de",
    "el",
    "en",
    "es",
    "et",
    "fa",
    "fi",
    "fr",
    "gu",
    "he",
    "hi",
    "hr",
    "hu",
    "hy",
    "id",
    "is",
    "it",
    "ja",
    "ka",
    "kk",
    "km",
    "kn",
    "ko",
    "lt",
    "lv",
    "mk",
    "ml",
    "mn",
    "mr",
    "my",
    "ne",
    "nl",
    "no",
    "pl",
    "pt",
    "ro",
    "ru",
    "uk",
    "zh",
]
PERPLEXITY_MODELS = ["Wikipedia", "OSCAR"]


class ContextLogger:
    def __init__(self, text: str = ""):
        self.text = text
        self.start_time = time.time()

    def __enter__(self):
        logger.info(self.text)

    def __exit__(self, type, value, traceback):
        logger.info(f"Took: {time.time() - self.start_time:.4f} seconds")


def generate_plot(
    df: pd.DataFrame,
    text_column: str,
    label_column: str,
    sample: Optional[int],
    dimensionality_reduction_function: Callable,
    model: SentenceTransformer,
    seed: int = 0,
    context_logger: Union[st.spinner, ContextLogger] = ContextLogger,
    hub_dataset: str = "",
) -> Tuple[Figure, Optional[Figure]]:
    if text_column not in df.columns:
        raise ValueError(
            f"The specified column name doesn't exist. Columns available: {df.columns.values}"
        )
    if label_column not in df.columns:
        df[label_column] = 0
    df = df.dropna(subset=[text_column, label_column])
    if sample:
        df = df.sample(min(sample, df.shape[0]), random_state=seed)
    with context_logger(text="Embedding text..."):
        embeddings = embed_text(df[text_column].values.tolist(), model)
    logger.info("Encoding labels")
    encoded_labels = encode_labels(df[label_column])
    with context_logger("Reducing dimensionality..."):
        embeddings_2d = dimensionality_reduction_function(embeddings)
    logger.info("Generating figure")
    hover_data = {
        text_column: df[text_column].values,
        label_column: encoded_labels.values,
    }
    # Round perplexity values
    values = df[label_column].values.round().astype(int)
    plot = draw_interactive_scatter_plot(
        hover_data,
        embeddings_2d[:, 0],
        embeddings_2d[:, 1],
        values,
    )
    # Special case for the registry dataset
    plot_registry = None
    if hub_dataset == REGISTRY_DATASET:
        encoded_labels = encode_labels(df["label"])
        hover_data = {
            text_column: df[text_column].values,
            "label": df["label"].values,
            label_column: df[label_column].values,
        }
        plot_registry = draw_interactive_scatter_plot(
            hover_data,
            embeddings_2d[:, 0],
            embeddings_2d[:, 1],
            encoded_labels.values,
            palette=Turbo256,
        )
    return plot, plot_registry