"""Aspect-based sentiment analysis predictive task."""

from __future__ import annotations

from collections.abc import Callable, Iterable, Sequence
from pathlib import Path
from typing import Any, override

import datasets
import dspy
import pydantic
import sklearn

from sieves.data import Doc
from sieves.model_wrappers import ModelType
from sieves.model_wrappers.types import ModelSettings
from sieves.serialization import Config
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.schemas.sentiment_analysis import (
    FewshotExample,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)
from sieves.tasks.predictive.sentiment_analysis.bridges import (
    DSPySentimentAnalysis,
    PydanticSentimentAnalysis,
)

_TaskBridge = DSPySentimentAnalysis | PydanticSentimentAnalysis


class SentimentAnalysis(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Estimate per‑aspect and overall sentiment for a document."""

    THRESHOLD: float = 0.5

    def __init__(
        self,
        model: TaskModel,
        model_settings: ModelSettings = ModelSettings(),
        aspects: tuple[str, ...] = tuple(),
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """
        Initialize SentimentAnalysis task.

        :param model: Model to use.
        :param model_settings: Settings for structured generation.
        :param aspects: Aspects to consider in sentiment analysis. Overall sentiment will always be determined. If
            empty, only overall sentiment will be determined.
        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        :param batch_size: Batch size to use for inference. Use -1 to process all documents at once.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param fewshot_examples: Few-shot examples.
        :param condition: Optional callable that determines whether to process each document.
        """
        self._aspects = tuple(sorted(set(aspects) | {"overall"}))
        super().__init__(
            model=model,
            model_settings=model_settings,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=False,
            prompt_instructions=prompt_instructions,
            fewshot_examples=fewshot_examples,
            condition=condition,
        )

    @property
    @override
    def fewshot_example_type(self) -> type[FewshotExample]:
        """Return few-shot example type.

        :return: Few-shot example type.
        """
        return FewshotExample

    @property
    @override
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        """Return the unified Pydantic prompt signature for this task.

        :return: Unified Pydantic prompt signature.
        """
        fields = {
            aspect: (
                float,
                pydantic.Field(
                    ...,
                    description=f"Sentiment score for the '{aspect}' aspect, ranging from 0 (Negative) to 1 "
                    f"(Positive).",
                    ge=0,
                    le=1,
                ),
            )
            for aspect in self._aspects
        }
        fields["score"] = (
            float | None,
            pydantic.Field(
                default=None,
                description="Provide an overall confidence score for the sentiment analysis, ranging from 0 to 1.",
            ),
        )

        return pydantic.create_model(
            "SentimentAnalysisOutput",
            __base__=pydantic.BaseModel,
            __doc__="Result of aspect-based sentiment analysis.",
            **fields,
        )  # type: ignore[no-matching-overload]

    @property
    @override
    def metric(self) -> str:
        return "F1 (Macro)"

    @override
    def _compute_metrics(self, truths: list[Any], preds: list[Any], judge: dspy.LM | None = None) -> dict[str, float]:
        """Compute corpus-level metrics.

        Computes F1 per aspect, then averages over aspects. Sentiments are discretized to 0 (Negative) or 1 (Positive)
        for F1.

        :param truths: List of ground truths.
        :param preds: List of predictions.
        :param judge: Optional DSPy LM instance to use as judge for generative tasks.
        :return: Dictionary of metrics.
        """
        aspects_list = list(self._aspects)
        # We will compute F1 per aspect, then average.
        # Sentiments are 0.0 to 1.0. We discretize them to 0 (Negative) or 1 (Positive) for F1.

        scores: list[float] = []

        for aspect in aspects_list:
            y_true: list[int] = []
            y_pred: list[int] = []

            for gold, pred in zip(truths, preds):
                if gold is not None:
                    assert isinstance(gold, TaskResult)
                    gold_val = gold.sentiment_per_aspect.get(aspect, 0.0)
                    y_true.append(1 if gold_val >= self.THRESHOLD else 0)
                else:
                    y_true.append(-1)

                if pred is not None:
                    assert isinstance(pred, TaskResult)
                    pred_val = pred.sentiment_per_aspect.get(aspect, 0.0)
                    y_pred.append(1 if pred_val >= self.THRESHOLD else 0)
                else:
                    y_pred.append(-1)

            # Compute Macro F1 for this aspect.
            if y_true:
                scores.append(sklearn.metrics.f1_score(y_true, y_pred, average="macro", zero_division=0))

        avg_score = sum(scores) / len(scores) if scores else 0.0

        return {self.metric: avg_score}

    @override
    def _init_bridge(self, model_type: ModelType) -> _TaskBridge:
        bridge_types: dict[ModelType, type[DSPySentimentAnalysis | PydanticSentimentAnalysis]] = {
            ModelType.dspy: DSPySentimentAnalysis,
            ModelType.outlines: PydanticSentimentAnalysis,
            ModelType.langchain: PydanticSentimentAnalysis,
        }

        try:
            bridge_type = bridge_types[model_type]

            return bridge_type(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                aspects=self._aspects,
                model_settings=self._model_settings,
                prompt_signature=self.prompt_signature,
                model_type=model_type,
                fewshot_examples=self._fewshot_examples,
            )
        except KeyError as err:
            raise KeyError(f"Model type {model_type} is not supported by {self.__class__.__name__}.") from err

    @staticmethod
    @override
    def supports() -> set[ModelType]:
        return {
            ModelType.dspy,
            ModelType.langchain,
            ModelType.outlines,
        }

    @override
    def _validate_fewshot_examples(self) -> None:
        for fs_example in self._fewshot_examples:
            assert isinstance(fs_example, FewshotExample)
            if any([aspect not in self._aspects for aspect in fs_example.sentiment_per_aspect]) or not all(
                [label in fs_example.sentiment_per_aspect for label in self._aspects]
            ):
                raise ValueError(
                    f"Aspect mismatch: {self._task_id} has aspects {self._aspects}. Few-shot examples have "
                    f"aspects {fs_example.sentiment_per_aspect.keys()}."
                )

    @property
    @override
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "aspects": self._aspects,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> datasets.Dataset:
        threshold = threshold or self.THRESHOLD

        # Define metadata.
        features = datasets.Features(
            {
                "text": datasets.Value("string"),
                "aspect": datasets.Sequence(datasets.Value("float32")),
                "score": datasets.Value("float32"),
            }
        )
        info = datasets.DatasetInfo(
            description=f"Aspect-based sentiment analysis dataset with aspects {self._aspects}. Generated with sieves "
            f"v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        aspects = self._aspects
        try:
            data = [
                (
                    doc.text,
                    doc.results[self._task_id].sentiment_per_aspect,
                    doc.results[self._task_id].score,
                )
                for doc in docs
            ]
        except KeyError as err:
            raise KeyError(f"Not all documents have results for this task with ID {self._task_id}") from err

        def generate_data() -> Iterable[dict[str, Any]]:
            """Yield results as dicts.

            :return: Results as dicts.
            """
            for text, scores, score in data:
                yield {"text": text, "aspect": [scores[aspect] for aspect in aspects], "score": score}

        # Create dataset.
        return datasets.Dataset.from_generator(generate_data, features=features, info=info)

    @override
    def distill(
        self,
        base_model_id: str,
        framework: DistillationFramework,
        data: datasets.Dataset | Sequence[Doc],
        output_path: Path | str,
        val_frac: float,
        init_kwargs: dict[str, Any] | None = None,
        train_kwargs: dict[str, Any] | None = None,
        seed: int | None = None,
    ) -> None:
        raise NotImplementedError

    @override
    def _evaluate_dspy_example(self, truth: dspy.Example, pred: dspy.Prediction, trace: Any, model: dspy.LM) -> float:
        # Compute per-aspect accuracy as 1 - abs(true_sentiment - pred_sentiment)
        # Average across all aspects (same approach as multi-label classification)
        accuracy = 0
        for aspect, sentiment in truth["sentiment_per_aspect"].items():
            if aspect in pred["sentiment_per_aspect"]:
                pred_sentiment = max(min(pred["sentiment_per_aspect"][aspect], 1), 0)
                accuracy += 1 - abs(sentiment - pred_sentiment)

        base_accuracy = accuracy / len(truth["sentiment_per_aspect"])

        # If score is available in both truth and pred, incorporate it into the metric.
        if "score" in truth and truth["score"] is not None and "score" in pred and pred["score"] is not None:
            score_accuracy = 1 - abs(truth["score"] - max(min(pred["score"], 1), 0))
            return (base_accuracy + score_accuracy) / 2

        return base_accuracy
