from collections import Counter
import numpy as np
import torch
import torch.nn.functional as F
from sklearn.metrics import f1_score
from torch import nn
from torch.utils.data import DataLoader
from transformers import Trainer
from src.models.dataset import EntitySentimentDataset


def compute_class_weights(examples: list[dict], n_classes: int) -> torch.Tensor:
    counts = Counter(e["label"] for e in examples)
    total = sum(counts.values())
    weights = [total / (n_classes * counts.get(i, 1)) for i in range(n_classes)]
    return torch.tensor(weights, dtype=torch.float)


def focal_loss(
    logits: torch.Tensor,
    labels: torch.Tensor,
    weight: torch.Tensor,
    gamma: float = 2.0,
) -> torch.Tensor:
    ce = F.cross_entropy(logits, labels, weight=weight, reduction="none")
    probs = F.softmax(logits, dim=-1)
    pt = probs.gather(1, labels.unsqueeze(1)).squeeze(1)
    return ((1 - pt) ** gamma * ce).mean()


class WeightedLossTrainer(Trainer):

    def __init__(self, *args, class_weights: torch.Tensor, loss_fn: str = "cross_entropy", focal_gamma: float = 2.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.class_weights = class_weights
        self.loss_fn = loss_fn
        self.focal_gamma = focal_gamma

    def compute_loss(self, model, inputs, return_outputs: bool = False, **kwargs):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        w = self.class_weights.to(outputs.logits.device)
        if self.loss_fn == "focal":
            loss = focal_loss(outputs.logits, labels, weight=w, gamma=self.focal_gamma)
        else:
            loss = nn.CrossEntropyLoss(weight=w)(outputs.logits, labels)
        return (loss, outputs) if return_outputs else loss


def reconstruct_triplets(
    yes_probs: np.ndarray, bin_labels: np.ndarray
) -> tuple[list[int], list[int]]:
    """Group consecutive (neg, neu, pos) triplets and take argmax."""
    preds3, labels3 = [], []
    for i in range(0, len(yes_probs) - 2, 3):
        preds3.append(int(np.argmax(yes_probs[i: i + 3])))
        labels3.append(int(np.argmax(bin_labels[i: i + 3])))
    return preds3, labels3


def make_compute_metrics(mode: str):
    if mode in ("marker", "qa_m"):
        def compute_metrics(eval_pred):
            logits, labels = eval_pred
            preds = np.argmax(logits, axis=-1)
            macro_f1 = f1_score(labels, preds, average="macro")
            per_class = f1_score(labels, preds, average=None, labels=[0, 1, 2])
            return {
                "macro_f1": macro_f1,
                "f1_negative": per_class[0],
                "f1_neutral": per_class[1],
                "f1_positive": per_class[2],
            }
    else:
        def compute_metrics(eval_pred):
            logits, labels = eval_pred
            preds = np.argmax(logits, axis=-1)
            bin_acc = float((preds == labels).mean())
            bin_f1 = float(f1_score(labels, preds, average="binary", pos_label=1))

            yes_probs = F.softmax(
                torch.tensor(logits, dtype=torch.float), dim=-1
            )[:, 1].numpy()

            preds3, labels3 = reconstruct_triplets(yes_probs, labels)

            macro_f1 = float(f1_score(preds3, labels3, average="macro")) \
                if preds3 else 0.0
            return {
                "macro_f1": macro_f1,
                "bin_accuracy": bin_acc,
                "bin_f1_yes": bin_f1,
            }

    return compute_metrics


def evaluate_qa_b_test(
    model,
    tokenizer,
    test_exs: list[dict],
    max_len: int,
    batch_size: int,
    device: torch.device,
) -> tuple[float, list[int], list[int]]:
    ds = EntitySentimentDataset(test_exs, tokenizer, max_len)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)

    all_yes_probs, all_bin_labels = [], []
    model.eval()
    with torch.no_grad():
        for batch in loader:
            logits = model(
                input_ids=batch["input_ids"].to(device),
                attention_mask=batch["attention_mask"].to(device),
            ).logits
            all_yes_probs.extend(
                F.softmax(logits, dim=-1)[:, 1].cpu().tolist()
            )
            all_bin_labels.extend(batch["labels"].tolist())

    preds3, labels3 = reconstruct_triplets(
        np.array(all_yes_probs), np.array(all_bin_labels)
    )

    macro_f1 = f1_score(labels3, preds3, average="macro")
    return macro_f1, preds3, labels3