# -*- coding: utf-8 -*-
# Copyright (c) Louis Brulé Naudet. All Rights Reserved.
# This software may be used and distributed according to the terms of License Agreement.
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import sys

from typing import (
    IO,
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Type,
    Tuple,
    Union,
    Mapping,
    TypeVar,
    Callable,
    Optional,
    Sequence,
)

import datasets
import faiss
import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
import torch

from datasets import load_dataset, Dataset, DatasetDict
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
from sklearn.decomposition import IncrementalPCA
from tqdm import tqdm
from umap import UMAP

from ragoon._logger import (
    Logger,
    TqdmToLogger
)

from ragoon._decorators import (
    memory,
    timer
)

logger = Logger()


class EmbeddingsDataLoader:
    """
    A class to load and process datasets to add embeddings using specified models.

    This class handles loading a dataset from Hugging Face, processing it to add embeddings
    using specified models, and provides methods to save and upload the processed dataset.

    Attributes
    ----------
    dataset_name : str
        The name of the dataset to load from Hugging Face.
   
    token : str
        The token for accessing Hugging Face API.
   
    model_configs : list of dict
        The list of dictionaries with model configurations to use for generating embeddings.
   
    batch_size : int
        The number of samples to process in each batch.
   
    dataset : datasets.DatasetDict
        The loaded and processed dataset.
   
    convert_to_tensor : bool, optional
        Whether the output should be one large tensor. Default is False.
   
    cuda_available : bool
        Whether CUDA is available for GPU acceleration.
   
    device : str, optional
        The device used for embedding processing if torch.cuda.is_available() is not reliable.
        Useful when using the Zero GPU on Hugging Face Space. Default is None.
   
    models : dict
        A dictionary to store loaded models and their configurations.

    Methods
    -------
    __init__(token, model_configs, dataset_name=None, dataset=None, batch_size=16, convert_to_tensor=False, device=None)
        Initializes the EmbeddingDatasetLoader with the specified parameters.
    
    load_dataset()
        Load the dataset from Hugging Face.
    
    load_model(model_name)
        Load the specified model.
    
    load_models()
        Load all specified models.
    
    encode(texts, model, query_prefix=None, passage_prefix=None)
        Create embeddings for a list of texts using a loaded model with optional prefixes.
    
    embed(batch, model, model_name, column="text", query_prefix=None, passage_prefix=None)
        Add embeddings columns to the dataset for each model.
    
    batch_embed(text)
        Embed a single text using all loaded models and return the results as a JSON string.
    
    process_splits(splits=None, column="text", load_all_models=True)
        Process specified splits of the dataset and add embeddings for each model.
    
    get_dataset()
        Return the processed dataset.
    
    save_dataset(output_dir)
        Save the processed dataset to disk.
    
    upload_dataset(repo_id)
        Upload the processed dataset to the Hugging Face Hub.
    """
    def __init__(
        self,
        model_configs: List[Dict[str, str]],
        token: str = None,
        dataset_name: Optional[str] = None,
        dataset: Optional[Union[Dataset, DatasetDict]] = None,
        batch_size: Optional[int] = 8,
        convert_to_tensor: Optional[bool] = False,
        device: Optional[str] = "cuda",
    ):
        """
        Initialize the EmbeddingDatasetLoader with the specified parameters.

        Parameters
        ----------
        model_configs : list of dict
            The list of dictionaries with model configurations to use for generating embeddings.

        token : str
            The token for accessing Hugging Face API. Default is None.

        dataset_name : str, optional
            The name of the dataset to load from Hugging Face. Default is None.

        dataset : Dataset or DatasetDict, optional
            The dataset to process. Default is None.

        batch_size : int, optional
            The number of samples to process in each batch. Default is 16.

        convert_to_tensor : bool, optional
            Whether the output should be one large tensor. Default is False.

        device : str, optional
            The device used for embedding processing if torch.cuda.is_available() is not reliable.
            Useful when using the Zero GPU on Hugging Face Space. Default is 'cuda'.
        """
        if token and not isinstance(token, str):
            raise ValueError("Invalid token. Please provide a non-empty string.")

        if not model_configs or not isinstance(model_configs, list):
            raise ValueError(
                "Invalid model configurations. Please provide a non-empty list of dictionaries."
            )
            
        if not isinstance(batch_size, int) or batch_size <= 0:
            raise ValueError("Invalid batch size. Please provide a positive integer.")

        self.dataset_name: str = dataset_name
        self.token: str = token
        self.model_configs: List[Dict[str, str]] = model_configs
        self.batch_size: int = batch_size
        self.dataset: Union[Dataset, DatasetDict] = dataset
        self.convert_to_tensor: bool = convert_to_tensor
        self.cuda_available: bool = torch.cuda.is_available()
        self.device: Optional[str] = device
        self.models: Dict[str, Dict] = {}
        self.active_config: Dict[str, Dict] = {}
        self.column: Optional[str] = None


    def load_dataset(self):
        """
        Load the dataset from Hugging Face.

        Raises
        ------
        Exception
            If the dataset fails to load from Hugging Face.
        """
        try:
            self.dataset = load_dataset(self.dataset_name)
            logger.info(f"Dataset '{self.dataset_name}' loaded successfully.")
        
        except Exception as e:
            logger.error(f"Failed to load dataset '{self.dataset_name}': {e}")
            self.dataset = None


    def load_model(
        self, model_name: str
    ) -> Union[InferenceClient, SentenceTransformer]:
        """
        Load the specified model.

        Parameters
        ----------
        model_name : str
            The name of the model to load.

        Returns
        -------
        model : Union[InferenceClient, SentenceTransformer]
            The loaded model, either a SentenceTransformer or an InferenceClient.

        Raises
        ------
        Exception
            If the model fails to load.
        """
        try:
            if self.cuda_available or self.device == "cuda":
                logger.info(f"Loading model {model_name} with CUDA.")
                return SentenceTransformer(model_name, trust_remote_code=True).to(
                    "cuda"
                )

            else:
                logger.info(
                    f"Loading model {model_name} using Hugging Face Inference API."
                )
                
                return InferenceClient(
                    token=self.token, 
                    model=model_name
                )

        except Exception as e:
            logger.error(f"Failed to load model '{model_name}': {e}")
            
            return None


    def load_models(
        self,
    ) -> Dict[
        str, Dict[str, Union[InferenceClient, SentenceTransformer, Optional[str]]]
    ]:
        """
        Load all specified models.

        This method loads all models specified in the `model_configs` and returns them
        in a dictionary format.

        Returns
        -------
        models : dict
            A dictionary where each key is a model name and each value is a dictionary
            containing the model and any prefixes.

        Examples
        --------
        >>> loader = EmbeddingDatasetLoader(token="your_token", model_configs=[{"model": "bert-base-uncased"}])
        >>> models = loader.load_models()
        """
        for config in self.model_configs:
            model_name = config["model"]
            query_prefix = config.get("query_prefix")
            passage_prefix = config.get("passage_prefix")
            model = self.load_model(
                model_name=model_name
            )
            
            self.models[model_name] = {
                "model": model,
                "query_prefix": query_prefix,
                "passage_prefix": passage_prefix,
            }

        logger.info(f"Loaded {len(self.models)} models.")
        
        return self.models


    def delete_model(
        self, 
        model: Union[InferenceClient, SentenceTransformer]
    ):
        """
        Delete the specified model and clear GPU cache.

        Parameters
        ----------
        model : Union[InferenceClient, SentenceTransformer]
            The model to delete.

        Returns
        -------
        None
        """
        del model
        torch.cuda.empty_cache()


    def encode(
        self,
        texts: List[str],
    ) -> Union[np.ndarray, dict]:
        """
        Create embeddings for a list of texts using a loaded model with optional prefixes,
        and optionally embed them into a batch.

        Parameters
        ----------
        texts : list of str or dict
            The list of texts to encode or a batch of data from the dataset.

        Returns
        -------
        np.ndarray or dict
            The embeddings for the texts, or the batch with added embedding columns.

        Raises
        ------
        Exception
            If encoding or embedding fails.
        """
        try:
            if self.column:
                texts = texts[self.column]

            assert not (
                self.active_config["query_prefix"]
                and self.active_config["passage_prefix"]
            ), "Only one of query_prefix or passage_prefix should be set, or neither."

            if self.active_config["query_prefix"]:
                texts = [
                    f"{self.active_config['query_prefix']}{text}" for text in texts
                ]

            if self.active_config["passage_prefix"]:
                texts = [
                    f"{self.active_config['query_prefix']}{text}" for text in texts
                ]

            if self.cuda_available:
                embeddings = self.active_config["model"].encode(
                    sentences=texts, convert_to_tensor=self.convert_to_tensor
                )

            else:
                embeddings = self.active_config["model"].feature_extraction(text=texts)

            # Ensure the embeddings are in the correct format
            if isinstance(embeddings, torch.Tensor):
                embeddings = embeddings.tolist()

            if isinstance(embeddings, np.ndarray):
                embeddings = embeddings.tolist()

            if not isinstance(embeddings[0], list):
                embeddings = [embeddings]

            # If column is specified, return embeddings as part of a batch dictionary
            if self.column:
                return {
                    f"{self._get_simple_model_name(self.active_config['model_name'])}_embeddings": embeddings
                }

            return embeddings

        except Exception as e:
            logger.error(f"Failed to encode texts: {e}")
            return None


    def batch_encode(
        self, 
        text: str
    ) -> str:
        """
        Embed a single text using all loaded models and return the results as a JSON string.

        Parameters
        ----------
        text : str
            The text to embed.

        Returns
        -------
        str
            The JSON string containing the embeddings from all models.

        Raises
        ------
        Exception
            If embedding fails.
        """
        try:
            results = {}
            for model_name, model_data in self.models.items():
                self.active_config = {
                    "model": model_data["model"],
                    "model_name": model_name,
                    "query_prefix": model_data["query_prefix"],
                    "passage_prefix": model_data["passage_prefix"],
                }

                embeddings = self.encode(
                    texts=[text],  # Pass a list containing the single text
                )

                results[
                    f"{self._get_simple_model_name(model_name)}_embeddings"
                ] = embeddings[0]

            return results

        except Exception as e:
            logger.error(f"Failed to batch embed text: {e}")
            return json.dumps({})

    def process(
        self,
        splits: Optional[List[str]] = None,
        column: Optional[str] = "text",
        preload_models: Optional[bool] = False,
    ):
        """
        Process specified splits of the dataset and add embeddings for each model.

        Parameters
        ----------
        splits : list of str, optional
            The list of splits to process. Default is None.
        
        column : str, optional
            The name of the column containing the text to encode. Default is "text".
        
        preload_models : bool, optional
            Whether to load all models specified in the `model_configs`. Default is True.

        Returns
        -------
        None

        Raises
        ------
        Exception
            If processing fails.
        """
        if self.dataset is None:
            raise ValueError(
                "Dataset not loaded. Please load a dataset before processing."
            )

        if splits is None and isinstance(self.dataset, DatasetDict):
            splits = list(self.dataset.keys())

        self.column = column

        if preload_models:
            self.load_models()
            if isinstance(self.dataset, DatasetDict):
                for split in tqdm(splits, desc="Processing splits"):
                    for model_name, model_data in self.models.items():
                        self.active_config = {
                            "model": model_data["model"],
                            "model_name": model_name,
                            "query_prefix": model_data["query_prefix"],
                            "passage_prefix": model_data["passage_prefix"],
                        }

                        self.dataset[split] = self.dataset[split].map(
                            self.encode, batched=True, batch_size=self.batch_size
                        )

                        self.delete_model(model_data["model"])

            elif isinstance(self.dataset, Dataset):
                for model_name, model_data in self.models.items():
                    self.active_config = {
                        "model": model_data["model"],
                        "model_name": model_name,
                        "query_prefix": model_data["query_prefix"],
                        "passage_prefix": model_data["passage_prefix"],
                    }
                    self.dataset = self.dataset.map(
                        self.encode, batched=True, batch_size=self.batch_size
                    )

                    self.delete_model(model_data["model"])

        else:
            for config in self.model_configs:
                if isinstance(self.dataset, DatasetDict):
                    self.active_config = {
                        "model": self.load_model(model_name=config["model"]),
                        "model_name": config["model"],
                        "query_prefix": config.get("query_prefix"),
                        "passage_prefix": config.get("passage_prefix"),
                    }

                    if self.active_config["model"] is not None:
                        for split in tqdm(splits, desc="Processing splits"):
                            self.dataset[split] = self.dataset[split].map(
                                self.encode, batched=True, batch_size=self.batch_size
                            )

                elif isinstance(self.dataset, Dataset):
                    self.active_config = {
                        "model": self.load_model(model_name=config["model"]),
                        "model_name": config["model"],
                        "query_prefix": config.get("query_prefix"),
                        "passage_prefix": config.get("passage_prefix"),
                    }

                    if self.active_config["model"] is not None:
                        self.dataset = self.dataset.map(
                            self.encode, batched=True, batch_size=self.batch_size
                        )

                self.delete_model(self.active_config["model"])


    def get_dataset(
        self
    ) -> Union[Dataset, DatasetDict]:
        """
        Return the processed dataset.

        Returns
        -------
        dataset : Union[Dataset, DatasetDict]
            The processed dataset.
        """
        return self.dataset


    def save_dataset(
        self, 
        output_dir: str
    ):
        """
        Save the processed dataset to disk.

        Parameters
        ----------
        output_dir : str
            The directory to save the dataset.

        Raises
        ------
        Exception
            If saving fails.
        """
        try:
            self.dataset.save_to_disk(output_dir)
            logger.info(f"Dataset saved to {output_dir}.")
        
        except Exception as e:
            logger.error(f"Failed to save dataset to {output_dir}: {e}")


    def upload_dataset(
        self, 
        repo_id: str,
        token: Optional[str] = None,
        private: Optional[bool] = False
    ):
        """
        Upload the processed dataset to the Hugging Face Hub.

        Parameters
        ----------
        repo_id : str
            The repository ID to upload the dataset.

        token : str, optional
            An optional authentication token for the Hugging Face Hub. If no token is passed, will default to the token saved 
            locally when logging in with huggingface-cli login. Will raise an error if no token is passed and the user is not logged-in.

        private : bool, optional
            Whether the dataset repository should be set to private or not.
            Only affects repository creation: a repository that already exists will not be affected by that parameter.

        Raises
        ------
        Exception
            If uploading fails.
        """
        try:
            self.dataset.push_to_hub(
                repo_id,
                token=token,
                private=private
            )
            logger.info(
                f"Dataset uploaded to Hugging Face Hub with repo ID: {repo_id}."
            )

        except Exception as e:
            logger.error(
                f"Failed to upload dataset to Hugging Face Hub with repo ID: {repo_id}: {e}"
            )


    def _get_simple_model_name(
        self, 
        model_name: str
    ) -> str:
        """
        Get the simplified model name (last part after the slash).

        Parameters
        ----------
        model_name : str
            The full model name.

        Returns
        -------
        str
            The simplified model name.
        """
        return model_name.split("/")[-1]


class EmbeddingsVisualizer:
    """
    A class for Embedding Exploration Lab, visualizing high-dimensional embeddings in 3D space.

    This class provides functionality to load embeddings from a FAISS index,
    reduce their dimensionality using PCA and/or t-SNE, and visualize them
    in an interactive 3D plot.

    Parameters
    ----------
    index_path : str
        Path to the FAISS index file.

    dataset_path : str
        Path to the dataset containing labels.

    Attributes
    ----------
    index_path : str
        Path to the FAISS index file.

    dataset_path : str
        Path to the dataset containing labels.

    index : faiss.Index or None
        Loaded FAISS index.

    dataset : datasets.Dataset or None
        Loaded dataset containing labels.

    vectors : np.ndarray or None
        Extracted vectors from the FAISS index.

    reduced_vectors : np.ndarray or None
        Dimensionality-reduced vectors.

    labels : list of str or None
        Labels from the dataset.

    Methods
    -------
    load_index() -> 'EmbeddingsVisualizer':
        Load the FAISS index from the specified file path.

    load_dataset() -> 'EmbeddingsVisualizer':
        Load the dataset containing labels from the specified file path.

    extract_vectors() -> 'EmbeddingsVisualizer':
        Extract all vectors from the loaded FAISS index.

    reduce_dimensionality(
        method: str = "umap",
        pca_components: int = 50,
        final_components: int = 3,
        random_state: int = 42
    ) -> 'EmbeddingsVisualizer':
        Reduce dimensionality of the extracted vectors with dynamic progress tracking.

    plot_3d() -> None:
        Generate a 3D scatter plot of the reduced vectors with labels.

    Examples
    --------
    >>> visualizer = EmbeddingsVisualizer(index_path="path/to/index", dataset_path="path/to/dataset")
    >>> visualizer.visualize(
    ...    method="pca",
    ...    save_html=True,
    ...    html_file_name="embedding_visualization.html"
    ... )

    """
    def __init__(
        self, 
        index_path: str,
        dataset_path: str
    ):
        self.index_path: str = index_path
        self.dataset_path: str = dataset_path
        self.index = None
        self.dataset = None
        self.vectors = None
        self.reduced_vectors = None


    @memory(print_report=True)
    @timer(print_time=True)
    def load_index(
        self
    ):
        """
        Load the FAISS index from the specified file path.

        Returns
        -------
        self : EmbeddingsVisualizer
            The instance itself, allowing for method chaining.
        """
        self.index = faiss.read_index_binary(
            self.index_path
        )

        return self

    
    def load_dataset(
        self,
        column: str = "document"
    ):
        """
        Load the Dataset containing labels from the specified file path.

        Parameters
        ----------
        column : str, optional
            The column of the split corresponding to the embeddings stored in 
            the index. Default is 'document'.

        Returns
        -------
        self : datasets.Dataset
            The instance itself, allowing for method chaining.
        """
        self.dataset = datasets.load_from_disk(
            dataset_path=self.dataset_path
        )

        self.labels = self.dataset[column]

        return self


    @memory(print_report=True)
    @timer(print_time=True)
    def extract_vectors(
        self
    ):
        """
        Extract all vectors from the loaded FAISS index.

        This method should be called after `load_index()`.

        Returns
        -------
        self : EmbeddingsVisualizer
            The instance itself, allowing for method chaining.

        Raises
        ------
        ValueError
            If the index has not been loaded yet.

        RuntimeError
            If there's an issue with vector extraction.
        """
        global logger

        if self.index is None:
            raise ValueError("Index not loaded. Call load_index() first.")

        if isinstance(self.index, faiss.IndexBinaryFlat):
            # Handle binary index
            num_vectors = self.index.ntotal
            dimension = self.index.d
            code_size = dimension // 8  # Binary vectors are typically in bits, hence // 8 for bytes
            
            logger.info(f"Expected dimension (bits): {dimension}")
            logger.info(f"Index total vectors: {num_vectors}")
            logger.info(f"Index code size (bytes): {code_size}")

            # Binary vectors are stored in uint8 arrays
            self.vectors = np.empty(
                (num_vectors, code_size), 
                dtype=np.uint8
            )

            logger.info(f"Initialized binary vectors array with shape: {self.vectors.shape}")
            
            try:
                for i in tqdm(range(num_vectors)):
                    self.vectors[i] = self.index.reconstruct(i)
            
            except AssertionError:
                raise RuntimeError(
                    f"Dimension mismatch error. Expected dimension (bits): {dimension}, "
                    f"Index total vectors: {num_vectors}, "
                    f"Index code size (bytes): {code_size}. "
                    "The index might be corrupted or incompatible."
                )

            except AttributeError:
                raise RuntimeError(
                    "The loaded index doesn't support binary vector reconstruction."
                    "Make sure you're using an appropriate index type."
                )

            except Exception as e:
                raise RuntimeError(
                    f"Error extracting vectors: {str(e)}"
                )

        else:
            num_vectors = self.index.ntotal
            dimension = self.index.d
            code_size = getattr(self.index, "code_size", None)

            logger.info(f"Expected dimension: {dimension}")
            logger.info(f"Index total vectors: {num_vectors}")
            logger.info(f"Index code size: {code_size}")

            if code_size is not None and code_size != dimension:
                raise RuntimeError(
                    f"Dimension mismatch error. Expected dimension: {dimension}, "
                    f"Index total vectors: {num_vectors}, "
                    f"Index code size: {code_size}. "
                    "The index might be corrupted or incompatible."
                )

            self.vectors = np.empty(
                (num_vectors, dimension), 
                dtype=np.float32
            )

            logger.info(f"Initialized vectors array with shape: {self.vectors.shape}")
            
            try:
                self.index.reconstruct_n(
                    0, 
                    num_vectors, 
                    self.vectors
                )

            except AssertionError:
                raise RuntimeError(
                    f"Dimension mismatch error. Expected dimension: {dimension}, "
                    f"Index total vectors: {num_vectors}, "
                    f"Index code size: {code_size}. "
                    "The index might be corrupted or incompatible."
                )
            except AttributeError:
                raise RuntimeError(
                    "The loaded index doesn't support vector reconstruction."
                    "Make sure you're using an appropriate index type."
                )

            except Exception as e:
                raise RuntimeError(f"Error extracting vectors: {str(e)}")
                
            return self
        

    @memory(print_report=True)
    @timer(print_time=True)
    def reduce_dimensionality(
        self,
        method: str = "umap",
        pca_components: int = 50,
        final_components: int = 3,
        random_state: int = 42
    ):
        """
        Reduce dimensionality of the extracted vectors with dynamic progress tracking.
        
        Parameters
        ----------
        method : {'pca', 'umap', 'pca_umap'}, optional
            The method to use for dimensionality reduction, by default 'umap'.

            - pca : Principal Component Analysis (PCA) is a linear dimensionality reduction technique 
            that is commonly used to reduce the dimensionality of high-dimensional data. 
            It identifies the directions (principal components) in which the data varies the most 
            and projects the data onto these components, resulting in a lower-dimensional representation.

            - umap : Uniform Manifold Approximation and Projection (UMAP) is a non-linear dimensionality 
            reduction technique that is particularly well-suited for visualizing high-dimensional 
            data in lower-dimensional space. It preserves both local and global structure of the data 
            by constructing a low-dimensional representation that captures the underlying manifold 
            structure of the data.

            - pca_umap : PCA followed by UMAP is a two-step dimensionality reduction technique. 
            First, PCA is applied to reduce the dimensionality of the data. 
            Then, UMAP is applied to further reduce the dimensionality and capture the 
            non-linear structure of the data. This combination can be effective in 
            preserving both global and local structure of the data.
        
        pca_components : int, optional
            Number of components for PCA (used in 'pca' and 'pca_umap'), by default 50.
        
        final_components : int, optional
            Final number of components (3 for 3D visualization), by default 3.
        
        random_state : int, optional
            Random state for reproducibility, by default 42.
        
        Returns
        -------
        self : EmbeddingsVisualizer
            The instance itself, allowing for method chaining.
        
        Raises
        ------
        ValueError
            If vectors have not been extracted yet or if an invalid method is specified.
        """
        if self.vectors is None:
            raise ValueError("Vectors not extracted. Call extract_vectors() first.")

        if method == "pca":
            pca = IncrementalPCA(
                n_components=final_components, 
                batch_size=100
            )

            batches = np.array_split(
                self.vectors, 
                max(1, len(self.vectors) // 100)
            )
            
            with tqdm(total=len(batches), desc="PCA") as pbar:
                for batch in batches:
                    pca.partial_fit(batch)
                    pbar.update(1)
            
            self.reduced_vectors = pca.transform(self.vectors)

        elif method == "umap":
            tqdm_out = TqdmToLogger()
            old_stdout = sys.stdout
            sys.stdout = tqdm_out
            
            umap = UMAP(
                n_components=final_components, 
                random_state=random_state, 
                verbose=True
            )
            self.reduced_vectors = umap.fit_transform(self.vectors)
            
            sys.stdout = old_stdout
            tqdm_out.pbar.close()

        elif method == "pca_umap":
            # First, perform PCA
            pca = IncrementalPCA(
                n_components=pca_components, 
                batch_size=100
            )

            batches = np.array_split(
                self.vectors, max(1, len(self.vectors) // 100)
            )
            
            with tqdm(total=len(batches), desc="PCA") as pbar:
                for batch in batches:
                    pca.partial_fit(batch)
                    pbar.update(1)
            
            pca_result = pca.transform(self.vectors)
            
            # Then, perform UMAP
            tqdm_out = TqdmToLogger()
            old_stdout = sys.stdout
            sys.stdout = tqdm_out
            
            umap = UMAP(
                n_components=final_components, 
                random_state=random_state, 
                verbose=True
            )

            self.reduced_vectors = umap.fit_transform(pca_result)
            
            sys.stdout = old_stdout
            tqdm_out.pbar.close()

        else:
            raise ValueError("Invalid method. Choose 'pca', 'umap', or 'pca_umap'.")

        return self


    @memory(print_report=True)
    @timer(print_time=True)
    def create_plot(
        self, 
        title: str = "3D Visualization of Embeddings", 
        point_size: int = 3
    ) -> go.Figure:
        """
        Generate a 3D scatter plot of the reduced vectors with labels.

        Parameters
        ----------
        title : str, optional
            The title of the plot (default is '3D Visualization of Embeddings').

        point_size : int, optional
            The size of the markers in the scatter plot (default is 3).

        Returns
        -------
        go.Figure
            The generated 3D scatter plot.

        Raises
        ------
        ValueError
            If vectors have not been reduced yet.

        Notes
        -----
        This method requires the `plotly` library to be installed.

        Examples
        --------
        >>> visualizer = EmbeddingsVisualizer()
        >>> plot = visualizer.create_plot(title='My Embeddings', point_size=5)
        >>> plot.show()
        """
        if self.reduced_vectors is None:
            raise ValueError("Dimensionality not reduced. Call reduce_dimensionality() first.")

        fig = go.Figure()

        # Start with an empty scatter plot
        scatter = go.Scatter3d(
            x=[], y=[], z=[],
            mode="markers",
            text=[],
            textposition="top center",
            hovertemplate="%{text}",
            hoverinfo="text",
            marker=dict(
                size=point_size,
                color=[],
                colorscale="Viridis",
                opacity=0.8
            )
        )

        fig.add_trace(scatter)

        # Update layout for full-window size and remove axes
        fig.update_layout(
            title=dict(
                text=title,
                y=0.95,
                x=0.5,
                xanchor="center",
                yanchor="top"
            ),
            scene=dict(
                xaxis=dict(visible=False),
                yaxis=dict(visible=False),
                zaxis=dict(visible=False),
                bgcolor="rgba(0,0,0,0)"
            ),
            paper_bgcolor="rgba(0,0,0,0)",
            plot_bgcolor="rgba(0,0,0,0)",
            margin=dict(l=0, r=0, t=50, b=50),  # Added some bottom margin for the slider
            autosize=True
        )

        # Add slider
        steps = self._create_slider_steps()
        sliders = [
            dict(
                active=0,
                currentvalue={
                    "prefix": "Number of embeddings: ", 
                    "xanchor": "center"
                },
                pad={
                    "b": 10, 
                    "t": 10
                },
                len=0.9,  # 90% of the plot width
                x=0.05,   # Start at 5% from the left
                xanchor="left",
                y=0,
                yanchor="bottom",
                steps=steps
            )
        ]
        
        fig.update_layout(
            sliders=sliders
        )

        return fig


    def _create_slider_steps(
        self, 
        num_steps: int = 100
    ) -> list:
        """
        Create steps for the slider.

        Parameters:
        ----------
        num_steps : int, optional
            The number of steps to create for the slider. Default is 100.

        Returns:
        -------
        list
            A list of dictionaries representing the steps for the slider.

        Notes:
        ------
        This method creates steps for a slider that can be used to update the visualization of embeddings.

        The steps are created based on the number of vectors in the `reduced_vectors` attribute of the class.
        The `num_steps` parameter determines the number of steps to create. The default value is 100.

        Each step is represented by a dictionary with the following keys:
        - 'method': The update method to be called when the step is selected.
        - 'args': The arguments to be passed to the update method.
        - 'label': The label to be displayed for the step.

        The 'method' key should be set to "update" and the 'args' key should be a list of two dictionaries.
        The first dictionary contains the updated values for the x, y, z, text, and marker.color attributes of the visualization.
        The second dictionary contains the updated title for the visualization.

        The 'label' key should be set to a string representation of the step index.

        Example:
        --------
        >>> visualizer = EmbeddingsVisualizer()
        >>> steps = visualizer._create_slider_steps(num_steps=50)
        >>> print(steps)
        [{'method': 'update', 'args': [{'x': [array([1, 2, 3])], 'y': [array([4, 5, 6])], 'z': [array([7, 8, 9])], 'text': [['label1', 'label2', 'label3']], 'marker.color': [array([7, 8, 9])]}, {'title': 'Embeddings (showing 0 out of 10)'}], 'label': '0'}, ...]
        """
        num_vectors: int = len(
            self.reduced_vectors
        )

        step_size = max(1, num_vectors // num_steps)
        steps: list = []


        def _insert_line_breaks(
            text: str, 
            interval: int
        ) -> str:
            """
            Insert line breaks into the text at the specified interval.

            Parameters
            ----------
            text : str
                The input text.

            interval : int
                The interval at which line breaks should be inserted.

            Returns
            -------
            str
                The text with line breaks inserted.

            Examples
            --------
            >>> _insert_line_breaks("Hello, world!", 5)
            'Hello,<br> world!'
            >>> _insert_line_breaks("Lorem ipsum dolor sit amet, consectetur adipiscing elit.", 10)
            'Lorem ipsu<br>m dolor si<br>t amet, co<br>nsectetur<br> adipiscing<br> elit.'
            """
            return "<br>".join(text[i:i+interval] for i in range(0, len(text), interval))


        for i in tqdm(range(0, num_vectors + 1, step_size)):
            # Add HTML line breaks in the text for better formatting
            hover_texts = [_insert_line_breaks(label, 100) for label in self.labels[:i]]
            
            step = dict(
                method="update",
                args=[
                    {
                        "x": [self.reduced_vectors[:i, 0]],
                        "y": [self.reduced_vectors[:i, 1]],
                        "z": [self.reduced_vectors[:i, 2]],
                        "text": [hover_texts],
                        "marker.color": [self.reduced_vectors[:i, 2]]
                    },
                    {
                        "title": f"Embeddings (showing {i} out of {num_vectors})"
                    }
                ],
                label=str(i)
            )
            steps.append(step)

        return steps


    def visualize(
        self, 
        column: str,
        method: str = "tsne", 
        pca_components: int = 50, 
        final_components: int = 3,
        random_state: int = 42,
        title: str = "3D Visualization of Embeddings",
        point_size: int = 3,
        save_html: bool = False,
        html_file_name: str = "embedding_visualization.html"
    ):
        """
        Full pipeline: load index, extract vectors, reduce dimensionality, and visualize.

        Parameters
        ----------
        column: str
            The column of the split corresponding to the embeddings stored in the index.

        method : str, optional
            The dimensionality reduction method to use. Default is 'tsne'.

        pca_components : int, optional
            The number of components to keep when using PCA for dimensionality reduction. Default is 50.

        final_components : int, optional
            The number of final components to visualize. Default is 3.

        random_state : int, optional
            The random state for reproducibility. Default is 42.

        title : str, optional
            The title of the visualization plot. Default is '3D Visualization of Embeddings'.

        point_size : int, optional
            The size of the points in the visualization plot. Default is 3.

        save_html : bool, optional
            Whether to save the visualization as an HTML file. Default is False.

        html_file_name : str, optional
            The name of the HTML file to save. Default is 'embedding_visualization.html'.

        Returns
        -------
        None

        Raises
        ------
        None

        Examples
        --------
        >>> visualizer = EmbeddingsVisualizer()
        >>> visualizer.visualize(method='tsne', pca_components=50, final_components=3, random_state=42)
        """
        self.load_index()
        self.load_dataset(
            column=column
        )
        self.extract_vectors()
        self.reduce_dimensionality(
            method=method,
            pca_components=pca_components,
            final_components=final_components,
            random_state=random_state
        )

        fig = self.create_plot(
            title=title,
            point_size=point_size
        )

        if save_html:
            pio.write_html(
                fig, 
                file=html_file_name,
                auto_open=False
            )

            logger.info(f"Visualization saved as {html_file_name}")

        fig.show(
            config={
                "displayModeBar": False,
                "scrollZoom": True
            }
        )