import os

from datasets import load_dataset
from qdrant_client import QdrantClient, models
from qdrant_client.conversions.common_types import CollectionInfo
from dotenv import load_dotenv, find_dotenv


class MeasureRetrievalQuality:
    def __init__(self, dataset_path: str, collection_name: str, streaming: bool = True):
        _ = load_dotenv(find_dotenv())
        # path = "Qdrant/arxiv-titles-instructorxl-embeddings"
        dataset = load_dataset(
            path=dataset_path, split="train", streaming=True,
            token=os.environ.get('HF_TOKEN')
        )
        self.collection_name = collection_name or os.environ.get('COLLECTION_NAME')
        dataset_iterator = iter(dataset)
        self.train_dataset = [next(dataset_iterator) for _ in range(10000)]
        self.test_dataset = [next(dataset_iterator) for _ in range(1000)]
        self.client = QdrantClient(url=os.environ.get('DB_URL'), api_key=os.environ.get('DB_API_KEY'))

        self._upset_and_index()

    def _upset_and_index(self):

        if not self.client.collection_exists(collection_name=self.collection_name):
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=models.VectorParams(
                    size=768,  # Size of the embeddings generated by InstructorXL model
                    distance=models.Distance.COSINE,
                )
            )

        self.client.upload_points(  # upload_points is available as of qdrant-client v1.7.1
            collection_name=self.collection_name,
            points=[
                models.PointStruct(
                    id=item["id"],
                    vector=item["vector"],
                    payload=item,
                )
                for item in self.train_dataset
            ]
        )

        # Collection status is green, which means the indexing is finished
        while True:
            collection_info = self.client.get_collection(collection_name=self.collection_name)
            if collection_info.status == models.CollectionStatus.GREEN:
                break

    def compute_avg_precision_at_k(self, k: int):
        precisions = []
        for item in self.test_dataset:
            ann_result = self.client.query_points(
                collection_name=self.collection_name,
                query=item["vector"],
                limit=k,
            ).points

            knn_result = self.client.query_points(
                collection_name=self.collection_name,
                query=item["vector"],
                limit=k,
                search_params=models.SearchParams(
                    exact=True,  # Turns on the exact search mode
                ),
            ).points

            # We can calculate the precision@k by comparing the ids of the search results
            ann_ids = set(item.id for item in ann_result)
            knn_ids = set(item.id for item in knn_result)
            precision = len(ann_ids.intersection(knn_ids)) / k
            precisions.append(precision)

        return sum(precisions) / len(precisions)

    def tune_hnsw_configs(self):
        # Tweaking the HNSW parameters
        self.client.update_collection(
            collection_name=self.collection_name,
            hnsw_config=models.HnswConfigDiff(
                m=32,  # Increase the number of edges per node from the default 16 to 32
                ef_construct=200,  # Increase the number of neighbours from the default 100 to 200
            )
        )

        # Collection status is green, which means the indexing is finished
        while True:
            collection_info = self.client.get_collection(collection_name=self.collection_name)
            if collection_info.status == models.CollectionStatus.GREEN:
                break