from abc import ABC, abstractmethod
import faiss
from typing import Tuple, Protocol, Optional, Union, Any, List
import numpy as np
from pupil.db.config import FaissConf
from .config import NDArray2D, SimilarityType
from nptyping import NDArray
import math
from scipy import spatial

class VectorDB(Protocol):

    def search(self, query: NDArray2D, n_results :int = 4) -> Tuple[NDArray2D, NDArray2D]:
        """Search embeddings

        Args:
            query (NDArray2D): Vectors to search
            n_results (int, optional): Number of results per query. Defaults to 4.

        Returns:
            Tuple[NDArray2D, NDArray2D]: Return (Distances, indices)
        """
        ...

    def train_index(self, embeddings: NDArray2D) -> None:
        """Index the data 

        Args:
            embeddings (NDArray2D): _description_
        """
        ...

    def add(self, embeddings: NDArray2D) -> None:
        """Add embeddings into the database

        Args:
            embeddings (NDArray2D):
        """
        ...
    
    def __len__(self):
        ...

    def __getitem__(self, i):
        ...

class FaissVectorDB:
    def __init__(
        self, 
        similarity_metric: SimilarityType = FaissConf.similarity_type,
        nlist: Optional[int] = FaissConf.nlist, 
        nprobe:Optional[int] = FaissConf.nprobe
        )->None:

        self.similarity_metric = similarity_metric
        self.nlist = nlist
        self.nprobe = nprobe
        
    def build_index(self, embeddings: NDArray2D) -> None:

        size, dim = embeddings.shape
        if not self.nlist: self.nlist = min(4096, 8 * round(math.sqrt(size)))

        if size < 4 * 10000:
            fac_str = "Flat"
        elif size < 80 * 10000:
            fac_str = "IVF" + str(self.nlist) + ",Flat"
        elif size < 200 * 10000:
            fac_str = "IVF16384,Flat"
        else:
            fac_str = "IVF16384,PQ8"

            
        self.index = faiss.index_factory(dim, fac_str, faiss.METRIC_INNER_PRODUCT)
        self.index.nprobe = min(self.nprobe, self.nlist) # type: ignore

        if not self.index.is_trained:
            self.index.train(embeddings)
        if self.similarity_metric == SimilarityType.COSINE:
            faiss.normalize_L2(embeddings)
        self.add(embeddings)

    def __getitem__(self, i: Union[int, List[int], NDArray[(Any, ), np.int32]]):
        if not self.index.is_trained:
            raise ValueError("First add data to the database.")

        if type(i) is not int: 
            return np.vstack([self.__getitem__(ind) for ind in i]) # type: ignore

        if hasattr(self.index, 'make_direct_map'):
            self.index.make_direct_map()
        return self.index.reconstruct(i).reshape(1,-1).astype(np.float32)

    def add(self, emb: NDArray2D) -> None:
        self.index.add(emb)

    def search(self, query: NDArray2D, n_results :int = 4) -> Tuple[NDArray2D, NDArray2D]:
        if self.similarity_metric == SimilarityType.COSINE:
            faiss.normalize_L2(query)
        distances, inds = self.index.search(query, n_results + 1) 
        return distances[:, 1:], inds[:, 1:] # ((self.index.ntotal, n_resutls),  (self.index.ntotal, n_resutls))

    def __len__(self,):
        return self.index.ntotal