import itertools
from collections.abc import Iterable
from pathlib import Path
from typing import Any

from sieves.data.doc import Doc
from sieves.tasks.core import Task
from sieves.tasks.postprocessing.distillation.types import DistillationFramework, DistillationFrameworkLiteral
from sieves.tasks.predictive.core import PredictiveTask


class Distillation(Task):
    """Distills pipeline results into specialized models."""

    def __init__(
        self,
        target_task_id: str,
        base_model_id: str,
        framework: DistillationFramework | DistillationFrameworkLiteral,  # type: ignore[valid-type]
        output_path: Path | str,
        init_kwargs: dict[str, Any] | None = None,
        train_kwargs: dict[str, Any] | None = None,
        train_frac: float = 0.8,
        val_frac: float = 0.2,
        threshold: float = 0.5,
        seed: int | None = None,
        task_id: str | None = None,
        include_meta: bool = False,
    ):
        """Initializes a distillation task.

        :param target_task_id: ID of task whose results to distill.
        :param base_model_id: ID of Hugging Face model to use as base for distillation. The chosen model will be
            fine-tuned on the target task's results.
        :param framework: Which distillation framework to use.
        :param output_path: Path to store distilled model and training metadata at.
        :param init_kwargs: Kwargs passed on to model/trainer initialization.
        :param train_kwargs: Kwargs passed on to training call.
        :param train_frac: Fractions for training set. `train_frac` + `val_frac` must sum up to 1.
        :param val_frac: Fractions for validation set. `train_frac` + `val_frac` must sum up to 1.
        :param threshold: Threshold to apply when converting logits/confidence scores into labels or other structured
            predictions.
        :param seed: RNG seed.
        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        """
        super().__init__(task_id=task_id, include_meta=include_meta)

        self._target_task_id = target_task_id
        self._base_model_id = base_model_id
        self._framework = DistillationFramework(framework)
        self._output_path = Path(output_path)
        self._train_frac = train_frac
        self._val_frac = val_frac
        self._init_kwargs = init_kwargs or {}
        self._train_kwargs = train_kwargs or {}
        self._seed = seed
        self._threshold = threshold

        self._target_task: PredictiveTask | None = None  # type: ignore[type-arg]

    @property
    def target_task_id(self) -> str:
        """Get target task ID.
        :return: Target task ID.
        """
        return self._target_task_id

    @property
    def target_task(self) -> PredictiveTask | None:  # type: ignore[type-arg]
        """Get target task.
        :return: Target task.
        """
        return self._target_task

    @target_task.setter
    def target_task(self, value: PredictiveTask) -> None:  # type: ignore[type-arg]
        """Set target task.
        :param value: Target task.
        """
        self._target_task = value

    def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Distill results in docs for specified target task into local model. This will train and save a model in line
        with the arguments specified at task initialization time.

        :param docs: Docs with results to distill.
        :return: Unchanged documents.
        :raises ValueError: If `target_task_type` hasn't been set yet.
        """
        assert self._target_task, ValueError("`target_task` has to be set before running distillation.")
        assert isinstance(self._target_task, PredictiveTask), (
            "Only results of tasks of type `PredictiveTask` can be distilled."
        )
        docs_iter = itertools.tee(docs, 2)

        self._target_task.distill(
            base_model_id=self._base_model_id,
            distillation_framework=self._framework,
            hf_dataset=self._target_task.to_hf_dataset(docs_iter[0], threshold=self._threshold),
            output_path=self._output_path,
            init_kwargs=self._init_kwargs,
            train_kwargs=self._train_kwargs,
            train_frac=self._train_frac,
            val_frac=self._val_frac,
            seed=self._seed,
        )

        yield from docs_iter[1]

    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "target_task_id": self._target_task_id,
            "base_model_id": self._base_model_id,
            "framework": self._framework.value,
            "output_path": str(self._output_path),
            "init_kwargs": self._init_kwargs,
            "train_kwargs": self._train_kwargs,
            "train_frac": self._train_frac,
            "val_frac": self._val_frac,
            "threshold": self._threshold,
        }
