#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
This module provides classes and methods for managing vector databases using FAISS, generating embeddings using
OpenAI and Sentence-BERT models, and clustering documents using Gaussian Mixture Models (GMM) after dimensionality
reduction with UMAP. It supports adding documents to a FAISS index, searching for similar documents, and clustering
documents based on their embeddings.

Classes:
- BaseEmbeddingModel: An abstract base class for embedding models.
- OpenAIEmbedding: Embedding model using OpenAI's API.
- SentenceBERTEmbedding: A class for generating embeddings using the Sentence-BERT model.
- GMMCluster: Provides static methods for clustering embeddings using Gaussian Mixture Models.
- FAISSVectorDB: Manages a vector database using FAISS (Facebook AI Similarity Search).

Dependencies:
- logging: For logging messages.
- pickle: For saving and loading the state of the FAISS index and document metadata.
- typing: For type hints.
- faiss: For efficient similarity search and clustering of dense vectors.
- numpy: For numerical operations.
- pathlib: For handling filesystem paths.
- umap: For dimensionality reduction.
- langchain_openai: For generating embeddings using OpenAI's API.
- functools: For method overloading.
- sentence_transformers: For generating embeddings using the Sentence-BERT model.
- sklearn.mixture: For Gaussian Mixture Models clustering.
"""
import logging
import pickle
from typing import List, Tuple, Dict, Optional

import faiss
import numpy as np
from pathlib import Path

import umap
from langchain_openai import OpenAIEmbeddings
from functools import singledispatchmethod

from sentence_transformers import SentenceTransformer
from sklearn.mixture import GaussianMixture

logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)


class BaseEmbeddingModel(object):
    """
    A base class for embedding models.

    This class serves as an abstract base class for embedding models. It defines a common interface
    for embedding text into numerical vectors. Subclasses should implement the embedding method to
    convert text into a list of floats representing the embedding.

    Methods:
        embedding(text: str) -> List[float]:
            Abstract method for embedding text. Subclasses must override this method.
    """

    def embedding(self, text: str) -> List[float]:
        """
        Abstract method for embedding text.

        This method should be implemented by subclasses to convert a given piece of text into
        a list of floats representing the embedding.

        Parameters:
            text (str): The text to embed.

        Returns:
            List[float]: The embedding of the text as a list of floats.

        Raises:
            NotImplementedError: If the method is not overridden by a subclass.
        """
        raise NotImplementedError


class OpenAIEmbedding(BaseEmbeddingModel):
    """
    Embedding model using OpenAI's API.

    This class implements the embedding model using OpenAI's API. It initializes the model with a specified
    model name and dimensions for the embeddings.

    Attributes:
        model (OpenAIEmbeddings): The OpenAI embedding model instance.
    """
    def __init__(self, model: str, dimensions: int = 1024):
        """
        Initializes the OpenAI Embedding model with the specified model and dimensions.

        This constructor sets up an OpenAI Embedding model by specifying the model name and the dimensions
        of the embeddings it should produce. The model name is a string that identifies which OpenAI model
        to use for generating embeddings. The dimensions parameter specifies the size of the embedding vectors
        that the model will generate. By default, the dimension size is set to 1024.

        Parameters:
        - model (str): The name of the OpenAI model to use for generating embeddings.
        - dimensions (int, optional): The size of the embedding vectors. Defaults to 1024.

        Example:
        >>> embedding_model = OpenAIEmbedding(model="text-similarity-ada-001", dimensions=1024)
        """
        self.model = OpenAIEmbeddings(
            model=model,
            dimensions=dimensions
        )

    def embedding(self, text: str) -> List[float]:
        """
        Generates an embedding for the given text using the OpenAI model.

        This method takes a string of text as input and returns a list of floats representing the
        embedding generated by the OpenAI model. The embedding represents the text in a high-dimensional
        space, capturing semantic meaning that can be used for various natural language processing tasks.

        Parameters:
        - text (str): The text to generate an embedding for.

        Returns:
        - List[float]: A list of floats representing the generated embedding.

        Example:
        >>> model = OpenAIEmbedding(model="text-similarity-ada-001")
        >>> embedding = model.embedding("Hello, world!")
        >>> print(embedding)
        [0.1, 0.2, 0.3, ...]
        """
        return self.model.embed_query(text)


class SentenceBERTEmbedding(BaseEmbeddingModel):
    """
    A class for generating embeddings using the Sentence-BERT model.

    This class inherits from BaseEmbeddingModel and implements the embedding generation
    using the Sentence-BERT model. It is initialized with the path to the Sentence-BERT
    model and an optional maximum sequence length. The embedding method encodes the input
    text into a list of floats representing the high-dimensional vector.

    Attributes:
        model (SentenceTransformer): The Sentence-BERT model instance.
        model.max_seq_length (int): The maximum sequence length for the model. Defaults to 8192.

    Methods:
        embedding(text: str) -> List[float]:
            Generates an embedding for the given text using the Sentence-BERT model.
    """

    def __init__(self, model_path: str, max_seq_length: int = 8192):
        """
        Initializes the SentenceBERTEmbedding model with the specified model path and maximum sequence length.

        Parameters:
        - model_path (str): The file path to the pre-trained Sentence-BERT model.
        - max_seq_length (int, optional): The maximum sequence length that the model can handle. Defaults to 8192.

        Example:
        >>> sentence_bert_model = SentenceBERTEmbedding(model_path="path/to/sentence_bert_model")
        """
        self.model = SentenceTransformer(
            model_path,
            trust_remote_code=True
        )
        self.model.max_seq_length = max_seq_length

    def embedding(self, text: str) -> List[float]:
        """
        Generates an embedding for the given text using the Sentence-BERT model.

        This method takes a string of text as input and returns a list of floats representing the
        embedding generated by the Sentence-BERT model. The embedding represents the text in a
        high-dimensional space, capturing semantic meaning that can be used for various natural
        language processing tasks.

        Parameters:
        - text (str): The text to generate an embedding for.

        Returns:
        - List[float]: A list of floats representing the generated embedding.

        Example:
        >>> sentence_bert_model = SentenceBERTEmbedding(model_path="path/to/sentence_bert_model")
        >>> embedding = sentence_bert_model.embedding("Hello, world!")
        >>> print(embedding)
        [0.1, 0.2, 0.3, ...]
        """
        return self.model.encode([text])[0].tolist()


class GMMCluster(object):
    @staticmethod
    def global_cluster_embeddings(
            embeddings: np.ndarray,
            dim: int,
            n_neighbors: Optional[int] = None,
            metric: str = "cosine",
    ) -> np.ndarray:
        if n_neighbors is None:
            n_neighbors = int((len(embeddings) - 1) ** 0.5)
        reduced_embeddings = umap.UMAP(
            n_neighbors=n_neighbors, n_components=dim, metric=metric
        ).fit_transform(embeddings)
        return reduced_embeddings

    @staticmethod
    def local_cluster_embeddings(
            embeddings: np.ndarray, dim: int, num_neighbors: int = 10, metric: str = "cosine"
    ) -> np.ndarray:
        reduced_embeddings = umap.UMAP(
            n_neighbors=num_neighbors, n_components=dim, metric=metric
        ).fit_transform(embeddings)
        return reduced_embeddings

    @staticmethod
    def get_optimal_clusters(
            embeddings: np.ndarray, max_clusters: int = 50, random_state: int = 42
    ) -> int:
        max_clusters = min(max_clusters, len(embeddings))
        n_clusters = np.arange(1, max_clusters)
        bics = []
        for n in n_clusters:
            gm = GaussianMixture(n_components=n, random_state=random_state)
            gm.fit(embeddings)
            bics.append(gm.bic(embeddings))
        optimal_clusters = n_clusters[np.argmin(bics)]
        return optimal_clusters

    @staticmethod
    def GMM_cluster(embeddings: np.ndarray, threshold: float, random_state: int = 0):
        n_clusters = GMMCluster.get_optimal_clusters(embeddings)
        gm = GaussianMixture(n_components=n_clusters, random_state=random_state)
        gm.fit(embeddings)
        probs = gm.predict_proba(embeddings)
        labels = [np.where(prob > threshold)[0] for prob in probs]
        return labels, n_clusters

    @staticmethod
    def perform_clustering(
            embeddings: np.ndarray, dim: int, threshold: float, verbose: bool = False
    ) -> List[np.ndarray]:
        reduced_embeddings_global = GMMCluster.global_cluster_embeddings(embeddings, min(dim, len(embeddings) - 2))
        global_clusters, n_global_clusters = GMMCluster.GMM_cluster(
            reduced_embeddings_global, threshold
        )

        logging.info(f"Global Clusters: {n_global_clusters}")

        all_local_clusters = [np.array([]) for _ in range(len(embeddings))]
        total_clusters = 0

        for i in range(n_global_clusters):
            global_cluster_embeddings_ = embeddings[
                np.array([i in gc for gc in global_clusters])
            ]
            logging.info(
                f"Nodes in Global Cluster {i}: {len(global_cluster_embeddings_)}"
            )
            if len(global_cluster_embeddings_) == 0:
                continue
            if len(global_cluster_embeddings_) <= dim + 1:
                local_clusters = [np.array([0]) for _ in global_cluster_embeddings_]
                n_local_clusters = 1
            else:
                reduced_embeddings_local = GMMCluster.local_cluster_embeddings(
                    global_cluster_embeddings_, dim
                )
                local_clusters, n_local_clusters = GMMCluster.GMM_cluster(
                    reduced_embeddings_local, threshold
                )

            if verbose:
                logging.info(f"Local Clusters in Global Cluster {i}: {n_local_clusters}")

            for j in range(n_local_clusters):
                local_cluster_embeddings_ = global_cluster_embeddings_[
                    np.array([j in lc for lc in local_clusters])
                ]
                indices = np.where(
                    (embeddings == local_cluster_embeddings_[:, None]).all(-1)
                )[1]
                for idx in indices:
                    all_local_clusters[idx] = np.append(
                        all_local_clusters[idx], j + total_clusters
                    )

            total_clusters += n_local_clusters

        logging.info(f"Total Clusters: {total_clusters}")
        return all_local_clusters


class FAISSVectorDB(object):
    """
    A class for managing a vector database using FAISS (Facebook AI Similarity Search).

    This class provides functionality to initialize a FAISS index, add documents to the index,
    retrieve documents by their index, iterate over documents, and initialize the FAISS index
    with specific parameters.

    Attributes:
        embedding (BaseEmbeddingModel): The embedding model used to convert text to vectors.
        vector_db (faiss.Index): The FAISS index for storing vectors.
        documents (dict): A dictionary mapping document indices to their data.
        cursor (int): The current position for adding new documents.
        dimension (int): The dimensionality of the vectors stored in the index.
    """

    def __init__(self, embedding: BaseEmbeddingModel):
        """
        Initializes the FAISSVectorDB with a specified embedding model.

        Parameters:
            embedding (BaseEmbeddingModel): The embedding model to use for converting text to vectors.
        """
        self.embedding = embedding
        self.vector_db = None
        self.documents = {}
        self.cursor = 0
        self.dimension = None

    def __getitem__(self, idx: int) -> Dict:
        """
        Retrieves a document from the database by its index.

        Parameters:
            idx (int): The index of the document to retrieve.

        Returns:
            Dict: The document data.

        Raises:
            AssertionError: If the index is out of bounds.
        """
        assert idx < self.cursor, "Index out of bounds."
        return self.documents.get(idx)

    def __iter__(self):
        """
        Initializes the iterator by resetting the index counter.

        Returns:
            self: The instance itself to support iterator protocol.
        """
        self.idx = 0
        return self

    def __next__(self) -> Dict:
        """
        Retrieves the next document in the database.

        Returns:
            Dict: The next document data.

        Raises:
            StopIteration: If there are no more documents to iterate over.
        """
        if self.idx >= self.cursor:
            raise StopIteration
        document = self[self.idx]
        self.idx += 1
        return document

    def init_index(self, *args, force_init=False):
        """
        Initializes the FAISS index with optional parameters or defaults if not provided.

        If the index already exists and force_init is False, a ValueError is raised. If no
        parameters are provided, the index is initialized with the dimensionality of the embedding
        model's output and default settings.

        Parameters:
            *args: Variable length argument list for FAISS index_factory function.
            force_init (bool): If True, forces reinitialization of the index even if it already exists.

        Raises:
            ValueError: If the index already exists and force_init is False.
        """
        if self.vector_db and not force_init:
            raise ValueError("Index already exists.")
        if not args:
            dimensions = len(self.embedding.embedding("dummy"))
            args = (dimensions, "Flat", faiss.METRIC_INNER_PRODUCT)
        else:
            dimensions = args[0]
        self.vector_db = faiss.index_factory(*args)
        self.documents = {}
        self.cursor = 0
        self.dimension = dimensions

    @singledispatchmethod
    def add(self, data: Dict | List[Dict]):
        """
        Adds a single document or a list of documents to the FAISS index. This method is overloaded
        to accept either a dictionary representing a single document or a list of dictionaries,
        each representing a document. The actual implementation for each data type is provided
        by the registered methods below.

        Parameters:
        - data (Dict | List[Dict]): The document(s) to be added to the index. Each document must
          include a 'text' key with its content.

        Raises:
        - NotImplementedError: If the method is called without being overridden by a registered type.
        """

    @add.register(dict)
    def _(self, data: Dict):
        """
        Adds a single document to the FAISS index. The document is first converted to an embedding
        using the specified embedding model, then normalized and added to the FAISS index.

        Parameters:
        - data (Dict): The document to be added. Must contain a 'text' key.

        Raises:
        - AssertionError: If the document does not contain a 'text' key.
        - ValueError: If the FAISS index has not been initialized or loaded.
        """
        assert "text" in data, "Data must contain 'text' key."
        if not self.vector_db:
            raise ValueError("Index not initialized or loaded.")
        embedding = self.embedding.embedding(data.get("text"))
        vector = np.array(embedding).reshape(1, -1)
        vector /= np.linalg.norm(vector)
        self.vector_db.add(vector)
        self.documents[self.cursor] = {
            "vector": vector,
            "meta_data": data
        }
        self.cursor += 1

    @add.register(list)
    def _(self, data: List[Dict]):
        """
        Adds a list of documents to the FAISS index. Each document in the list is processed
        individually in the same manner as the single document add method.

        Parameters:
        - data (List[Dict]): A list of documents to be added. Each document must contain a 'text' key.
        """
        for d in data:
            self.add(d)

    def save(self, path: str | Path):
        """
        Saves the current state of the FAISS index and the associated document metadata to disk.

        Parameters:
        - path (str | Path): The directory path where the index and metadata should be saved.

        Raises:
        - ValueError: If the FAISS index has not been initialized or loaded.
        """
        if not self.vector_db:
            raise ValueError("Index not initialized or loaded.")
        if isinstance(path, str):
            path = Path(path)
        faiss.write_index(self.vector_db, str(path / "index.faiss"))
        with (path / "documents.pkl").open("wb") as f:
            state_dict = {
                "documents": self.documents,
                "cursor": self.cursor,
                "dimension": self.dimension
            }
            pickle.dump(state_dict, f)

    def load(self, path: str | Path, force_load=False):
        """
        Loads the FAISS index and document metadata from disk. If an index is already loaded,
        it can be forcibly reloaded using the force_load parameter.

        Parameters:
        - path (str | Path): The directory path from where the index and metadata should be loaded.
        - force_load (bool): If True, forces the index to be reloaded even if one is already loaded.

        Raises:
        - ValueError: If an index is already loaded and force_load is False.
        """
        if self.vector_db and not force_load:
            raise ValueError("Index already exists.")
        if isinstance(path, str):
            path = Path(path)
        self.vector_db = faiss.read_index(str(path / "index.faiss"))
        with (path / "documents.pkl").open("rb") as f:
            state_dict = pickle.load(f)
            self.documents = state_dict["documents"]
            self.cursor = state_dict["cursor"]
            self.dimension = state_dict["dimension"]

    def search(self, query: str, k: int = 10) -> List[Tuple[int, float, Dict]]:
        """
        Searches the FAISS index for the k most similar documents to the given query.

        This method first checks if the FAISS index is initialized and loaded. It then refines the
        number of results (k) based on the total documents in the index. The query is converted into
        an embedding, normalized, and used to search the FAISS index. The method returns a list of
        tuples, each containing the document index, the distance from the query, and the document data.

        Parameters:
        - query (str): The query string to search for.
        - k (int, optional): The number of similar documents to return. Defaults to 10.

        Returns:
        - List[Tuple[int, float, Dict]]: A list of tuples containing the index of the document in the
          FAISS index, the distance of the document from the query, and the document data.

        Raises:
        - ValueError: If the FAISS index is not initialized or loaded.
        """
        if not self.vector_db:
            raise ValueError("Index not initialized or loaded.")
        k_refine = max(min(k, self.cursor), 1)
        if k != k_refine:
            print(f"Refining k to {k_refine}.")
        query = self.embedding.embedding(query)
        query = np.array(query).reshape(1, -1)
        query /= np.linalg.norm(query)
        D, I = self.vector_db.search(query, k_refine)
        documents = [self[i] for i in I[0]]
        return list(zip(I[0], D[0], documents))

    def cluster(self, reduction_dimension: int = 10, threshold: float = 0.1) -> List[List[Dict]]:
        """
        Clusters the documents in the FAISS index using Gaussian Mixture Models (GMM) after dimensionality reduction.

        This method first checks if the FAISS index is initialized and loaded. It then performs dimensionality
        reduction on the document embeddings using UMAP, followed by clustering using GMM. The method returns
        a list of clusters, where each cluster is a list of documents that belong to that cluster.

        Parameters:
        - reduction_dimension (int, optional): The target dimensionality for the UMAP reduction. Defaults to 10.
        - threshold (float, optional): The threshold for assigning documents to clusters in GMM. Defaults to 0.1.

        Returns:
        - List[List[Dict]]: A list of clusters, each cluster being a list of documents that belong to it.

        Raises:
        - ValueError: If the FAISS index is not initialized or loaded.
        """
        if not self.vector_db:
            raise ValueError("Index not initialized or loaded.")
        dimension_refine = min(reduction_dimension, self.dimension)
        if reduction_dimension != dimension_refine:
            print(f"Refining reduction dimension to {dimension_refine}.")
        global_embeddings = np.array(list(i.get("vector") for i in self)).squeeze(axis=1)
        clusters = GMMCluster.perform_clustering(
            global_embeddings, dim=dimension_refine, threshold=threshold
        )
        node_clusters = []

        # Iterate over each unique label in the clusters
        for label in np.unique(np.concatenate(clusters)):
            # Get the indices of the nodes that belong to this cluster
            indices = [i for i, cluster in enumerate(clusters) if label in cluster]

            # Add the corresponding nodes to the node_clusters list
            cluster_nodes = [self[i] for i in indices]

            node_clusters.append(cluster_nodes)

        return node_clusters