"""
Implement operators to do vector-search and retrieval documents for a given user query. 
"""

from __future__ import annotations
import typing as t
import copy

import numpy as np
import polars as pl

if t.TYPE_CHECKING:
    from uptrain.framework import Settings
from uptrain.operators.base import (
    TransformOp,
    register_op,
    TYPE_TABLE_OUTPUT,
)
from uptrain.operators.embedding.embedding import Embedding
from uptrain.operators.io.base import JsonReader, CsvReader
from uptrain.utilities import lazy_load_dep, polars_to_pandas

faiss = lazy_load_dep("faiss", "faiss")


@register_op
class VectorSearch(TransformOp):
    """

     Attributes:
        col_in_query (str): Column name for the input query to be used for retrieval
        colk_in_document (str): Column name for the input document to be used for retrieval
        col_out (str): Column name for the retrieved_context
        documents (t.Union[list[str], str, JsonReader, CsvReader]): List of documents or path to the file containing the documents
        embeddings_model (str): Name of the embeddings model
        top_k (int): Top K documents will be retrieved.
        distance_metric (str): One of ['cosine_similarity', 'l2_distance']
        embedding_batch_size (int): Batch size for the embeddings model

    Raises:
        Exception: Raises exception for any failed evaluation attempts

    """

    col_in_query: str = "question"
    col_in_document: str = ""
    col_out: str = "context"
    documents: t.Union[list[str], str, JsonReader, CsvReader] = None
    embeddings_model: str = ""
    top_k: int = 1
    distance_metric: str = t.Literal["cosine_similarity", "l2_distance"]
    embedding_batch_size: int = 128

    def setup(self, settings: t.Optional[Settings] = None):
        read_op = None
        if isinstance(self.documents, list):
            documents_table = pl.DataFrame({"document": self.documents})
        elif isinstance(self.documents, str):
            if self.documents[-5:] == ".json":
                read_op = JsonReader(fpath=self.documents)
            elif self.documents[-6:] == ".jsonl":
                read_op = JsonReader(fpath=self.documents)
            elif self.documents[-4:] == ".csv":
                read_op = CsvReader(fpath=self.documents)
            else:
                read_op = JsonReader(fpath=self.documents)
                # raise Exception("File formats other than jsonl and csv are not supported")
        elif isinstance(self.documents, JsonReader) or isinstance(
            self.documents, CsvReader
        ):
            read_op = self.documents
        else:
            raise Exception(f"{type(self.documents)} is not supported")

        if read_op is not None:
            documents_table = pl.DataFrame(
                {
                    "document": read_op.setup(settings).run()["output"][
                        self.col_in_document
                    ]
                }
            )

        documents_table = pl.DataFrame(polars_to_pandas(documents_table))

        emb_op = Embedding(
            model=self.embeddings_model,
            col_in_text="document",
            col_out="document_embeddings",
            batch_size=self.embedding_batch_size,
        )
        doc_embeddings = np.array(
            list(
                emb_op.setup(settings).run(documents_table)["output"][
                    "document_embeddings"
                ]
            )
        )

        self.documents_list = documents_table["document"]

        if self.distance_metric == "cosine_similarity":
            self.vectorstore = faiss.IndexFlatIP(len(doc_embeddings[0]))
        elif self.distance_metric == "l2_distance":
            self.vectorstore = faiss.IndexFlatL2(len(doc_embeddings[0]))
        else:
            raise Exception(f"{self.distance_metric} is not allowed")
        self.vectorstore.add(doc_embeddings)

        self._settings = settings
        return self

    def run(self, data: pl.DataFrame) -> TYPE_TABLE_OUTPUT:
        emb_op = Embedding(
            model=self.embeddings_model,
            col_in_text=self.col_in_query,
            col_out=self.col_in_query + "_embeddings",
        )
        query_emb_res = emb_op.setup(self._settings).run(data)["output"].to_dicts()

        res = []
        for emb_res in query_emb_res:
            this_retrieved_scores, this_retrieved_idxs = self.vectorstore.search(
                np.array([emb_res[self.col_in_query + "_embeddings"]]), self.top_k
            )
            this_retrieved_documents = [
                self.documents_list[int(x)] for x in this_retrieved_idxs[0]
            ]

            del emb_res[self.col_in_query + "_embeddings"]

            this_retriev_res = copy.deepcopy(emb_res)
            this_retriev_res = []
            for _ in range(len(this_retrieved_documents)):
                this_retriev_res.append(copy.deepcopy(emb_res))
            for jdx, row in enumerate(this_retriev_res):
                row.update(
                    {
                        self.col_out: this_retrieved_documents[jdx],
                        "retrieval_rank": jdx + 1,
                        "retrieval_similarity_score": this_retrieved_scores[0][jdx],
                    }
                )
            res.extend(this_retriev_res)

        return {"output": pl.DataFrame(res)}
