from collections.abc import Sequence
from itertools import chain
from typing import Self, cast

import torch
from jaxtyping import Float
from torch import Tensor

from .types import Label


class Scores:
    """Object holding the scores on an unpaired dataset."""

    scores: list[Float[Tensor, " seq"]]  # of length num_dialogues
    labels: list[Label]  # of length num_dialogues

    def __init__(
        self, scores: Sequence[list[float] | torch.Tensor], labels: list[Label] | None = None
    ):
        self.scores = [torch.as_tensor(prompt_scores, dtype=torch.float32) for prompt_scores in scores]
        if labels is None:
            labels = [Label.HONEST] * len(self.scores)
        assert len(labels) == len(
            self.scores
        ), f"Mismatch between labels (len {len(labels)}) and scores (len {len(self.scores)})"
        self.labels = labels

    def to_dict(self) -> dict[str, list[list[float]] | list[str]]:
        return {
            "scores": [scores.tolist() for scores in self.scores],
            "labels": [label.name for label in self.labels],
        }

    @classmethod
    def from_dict(cls, dict: dict[str, list[list[float]] | list[str]]) -> Self:
        if "is_deceptive" in dict:  # to import legacy experiments
            labels = [
                Label.DECEPTIVE if is_deceptive else Label.HONEST
                for is_deceptive in dict["is_deceptive"]
            ]
        else:
            labels = [Label[cast(str, label)] for label in dict["labels"]]
        return cls(scores=[torch.tensor(scores) for scores in dict["scores"]], labels=labels)

    @classmethod
    def from_list(cls, scores_list: Sequence["Scores"]) -> "Scores":
        all_scores = list(chain.from_iterable([scores.scores for scores in scores_list]))
        all_labels = list(chain.from_iterable([scores.labels for scores in scores_list]))
        return cls(scores=all_scores, labels=all_labels)

    def __len__(self) -> int:
        return len(self.scores)

    def __getitem__(self, index: int) -> Float[Tensor, " seq"]:
        return self.scores[index]

    def get_scores_by_label(self) -> dict[Label, "Scores"]:
        def scores_for(target_label: Label) -> Scores:
            indices = [i for i, label in enumerate(self.labels) if label == target_label]
            return Scores([self.scores[i] for i in indices], [target_label] * len(indices))

        return {label: scores_for(label) for label in Label}
