from abc import ABC, abstractmethod
from typing import Dict, Any, Tuple
import torch
import numpy as np
from vox_bench.utils import Plotter, calculate_metrics, WandBHandler

class TaskHead(ABC):
    def __init__(self):
        self.best_metric = None

    @abstractmethod
    def forward(self, model, x_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        pass

    @abstractmethod
    def compute_loss(
        self,
        outputs: Dict[str, torch.Tensor],
        y_dict: Dict[str, torch.Tensor],
    ) -> torch.Tensor:
        pass

    @abstractmethod
    def evaluate_batch(
        self,
        outputs: Dict[str, torch.Tensor],
        y_dict: Dict[str, torch.Tensor],
    ) -> Dict[str, Any]:
        """
        returns values needed for metric aggregation
        """
        pass

    @abstractmethod
    def compute_epoch_metrics(
        self,
        aggregated: Dict[str, Any],
        prefix: str,
        epoch: int,
        wandb_handler,
    ) -> Dict[str, float]:
        pass

    @abstractmethod
    def get_monitor_value(self, metrics: Dict[str, float]) -> float:
        """
        balue used for scheduler and checkkpoint selection
        """
        pass

TASK_HEAD_REGISTRY : Dict[str, TaskHead] = {}
def register_head(name: str):
    def decorator(cls):
        TASK_HEAD_REGISTRY[name] = cls
        return cls
    return decorator

@register_head("classification")
class ClassificationHead(TaskHead):
    def __init__(
        self,
        criterion,
        class_names=None,
        threshold: float = 0.5,
    ):
        super().__init__()
        self.criterion = criterion
        self.class_names = class_names or ["Real", "Fake"]
        self.threshold = threshold

    def forward(self, model, x_dict):
        return model(**x_dict)

    def compute_loss(self, outputs, y_dict):
        return self.criterion(**(outputs | y_dict))

    def evaluate_batch(self, outputs, y_dict):
        logit = outputs["o1"].squeeze(1)
        target = y_dict["y1"]

        probs = torch.sigmoid(logit)
        preds = (probs > self.threshold).float()

        return {
            "y_true": target.detach().cpu(),
            "y_pred": preds.detach().cpu(),
            "y_probs": probs.detach().cpu(),
        }

    def compute_epoch_metrics(self, aggregated, prefix, epoch, wandb):

        y_true = torch.cat(aggregated["y_true"]).numpy()
        y_pred = torch.cat(aggregated["y_pred"]).numpy()
        y_probs = torch.cat(aggregated["y_probs"]).numpy()

        cm = Plotter.confusion_matrix(
            y_true, y_pred, self.class_names,
            prefix=prefix, epoch=epoch,
            wandb_handler=wandb,
        )

        scores = calculate_metrics(cm)

        roc_auc, roc_thr = Plotter.roc_auc_plt(
            y_true, y_probs,
            prefix=prefix, epoch=epoch,
            wandb_handler=self.wandb,
        )
        pr_auc, pr_thr = Plotter.pr_curve(
            y_true, y_probs,
            prefix=prefix, epoch=epoch,
            wandb_handler=self.wandb,
        )

        metrics = {
            f"{prefix}/loss":               avg_loss,
            f"{prefix}/f1":                 scores["f1 score"],
            f"{prefix}/recall":             scores["recall (fake detection rate)"],
            f"{prefix}/precision":          scores["precision (of all predicted fake, actually fake)"],
            f"{prefix}/far":                scores["false alarm rate (real predicted as fake)"],
            f"{prefix}/frr":                scores["false rejection rate (fake predicted as real)"],
            f"{prefix}/accuracy":           scores["accuracy"],
            f"{prefix}/roc_auc":            roc_auc,
            f"{prefix}/pr_auc":             pr_auc,
            f"{prefix}/roc_best_threshold": roc_thr,
            f"{prefix}/pr_best_threshold":  pr_thr,
        }

        self.wandb.log_metrics(metrics)
        return metrics

    def get_monitor_value(self, metrics):
        return metrics.get("val/f1", 0.0)
