import logging
import gzip
from pathlib import Path
from typing import Type
from experimaestro import Task, Param, Annotated, pathgenerator, Option, tqdm
import numpy as np
import datamaestro_text.data.ir as ir
from datamaestro_text.utils.shuffle import shuffle


def getpathname(context, config):
    name = "triplets.lst"
    if config.compressed:
        name = "triplets.lst.gz"

    return context.currentpath() / name


class StoreTrainingTripletTopicAdapter(ir.TrainingTriplets):
    """Retrieve an adhoc topic from a topic store and returns triplets with those"""

    id: Param[str] = ""

    store: Param[ir.TopicsStore]
    """The topic store to use"""

    data: Param[ir.TrainingTriplets]
    """Input data"""

    def __validate__(self):
        assert self.data.topic_cls.has_id, "Topics have no ID"

    def iter(self):
        for topic, doc1, doc2 in self.data.iter():
            yield self.store.topic_ext(topic.get_id()), doc1, doc2

    def count(self):
        return self.data.count()

    @property
    def topic_cls(self) -> Type[ir.Topic]:
        """The class for topics"""
        return self.store.topic_cls

    @property
    def document_cls(self) -> Type[ir.Document]:
        """The class for documents"""
        return self.data.document_cls


class StoreTrainingTripletDocumentAdapter(ir.TrainingTriplets):
    """Retrieve an adhoc topic from a topic store and returns triplets with those"""

    id: Param[str] = ""

    store: Param[ir.DocumentStore]
    """The topic store to use"""

    data: Param[ir.TrainingTriplets]
    """Input data"""

    def __validate__(self):
        assert self.data.document_cls.has_id, "Documents have no ID"

    def iter(self):
        for topic, doc1, doc2 in self.data.iter():
            yield topic, self.store.document_ext(
                doc1.get_id()
            ), self.store.document_ext(doc2.get_id())

    def count(self):
        return self.data.count()

    @property
    def topic_cls(self) -> Type[ir.Topic]:
        """The class for topics"""
        return self.store.topic_cls

    @property
    def document_cls(self) -> Type[ir.Document]:
        """The class for documents"""
        return self.data.document_cls


class ShuffledTrainingTripletsLines(Task):
    """Shuffle a set of training triplets"""

    data: Param[ir.TrainingTriplets]
    """Input data"""

    path: Annotated[Path, pathgenerator(getpathname)]
    """Output path"""

    doc_ids: Param[bool]
    """Whether to use document ids"""

    topic_ids: Param[bool]
    """True if we have query IDs"""

    seed: Param[int]
    """The random seed"""

    compressed: Option[bool] = True
    """Compress the output"""

    sample_rate: Param[float] = 1.0
    """Sampling rate - set to 1 to keep all the samples"""

    sample_max: Param[int] = 0
    """Maximum number of samples"""

    tmp_path: Annotated[Path, pathgenerator("tmp")]
    """Path where temporary files will be stored"""

    def __validate__(self):
        if self.topic_ids:
            assert self.data.topic_cls.has_id, "No topic ID in the source data"
        else:
            assert self.data.topic_cls.has_text, "No topic text in the source data"

        if self.doc_ids:
            assert self.data.document_cls.has_id, "No doc ID in the source data"
        else:
            assert self.data.document_cls.has_text, "No doc text in the source data"

    def task_outputs(self, dep):
        return dep(
            ir.TrainingTripletsLines(
                id="",
                path=self.path,
                topic_ids=self.topic_ids,
                doc_ids=self.doc_ids,
                sep="\t",
            )
        )

    def execute(self):
        # --- Shuffle using the shuf command with a seed

        random = np.random.RandomState(self.seed)

        if self.topic_ids:

            def get_query(query):
                return query.get_id()

        else:

            def get_query(query):
                return query.get_text()

        if self.doc_ids:

            def get_doc(doc):
                return doc.get_id()

        else:

            def get_doc(doc):
                return doc.get_text()

        def triplegenerator():
            logging.info("Starting to output triples")
            count = 0

            total = self.data.count()
            if self.sample_max > 0:
                total = min(total, self.sample_max)

            pbar = tqdm(total=total)
            for query, doca, docb in self.data.iter():
                # Discard sample
                if self.sample_rate < 1:
                    if random.uniform() > self.sample_rate:
                        continue

                pbar.update(1)
                count += 1
                yield f"{get_query(query)}\t{get_doc(doca)}\t{get_doc(docb)}\n"

                if self.sample_max > 0 and count >= self.sample_max:
                    break

            logging.info("Triples output ended (%d triples)", count)

        logging.info("Creating generator")

        # Output can be a stream or nothing
        if self.compressed:
            output = gzip.open(self.path, "wt")
        else:
            output = self.path.open("wt")

        with output:
            self.tmp_path.mkdir(exist_ok=True)
            shuffle(triplegenerator(), output, random=random, tmp_path=self.tmp_path)
