import torch
import numpy as np
from sklearn.metrics import (
    accuracy_score,
    roc_auc_score,
    f1_score,
    confusion_matrix,
)

# AUC and F1 are undefined when only one class is present in the batch;
# returns None for those fields rather than raising
def binary_metrics(logits: torch.Tensor, labels: torch.Tensor) -> dict:
    probs = torch.sigmoid(logits).numpy()
    preds = (probs >= 0.5).astype(int)
    y = labels.numpy().astype(int)
    has_both_classes = len(np.unique(y)) > 1
    auc_roc = float(roc_auc_score(y, probs)) if has_both_classes else None
    f1 = float(f1_score(y, preds, zero_division=0)) if has_both_classes else None
    return {
        "accuracy": float(accuracy_score(y, preds)),
        "auc_roc": auc_roc,
        "f1": f1,
        "confusion_matrix": confusion_matrix(y, preds, labels=[0, 1]),
    }


# Converts per-sample records to tensors and delegates to binary_metrics
def calc_metrics(records):
    logits = torch.tensor([r["logit"] for r in records], dtype=torch.float32)
    labels = torch.tensor([r["label"] for r in records], dtype=torch.float32)
    metrics = binary_metrics(logits, labels)
    metrics["confusion_matrix"] = metrics["confusion_matrix"].tolist()
    return metrics


# Per-source summaries; fake sources get detection_rate + pairwise_auc since AUC is undefined single-class
def source_metrics(records, real_source="wiki"):
    wiki_records = [r for r in records if r["source"] == real_source]
    by_source = {}
    for source in sorted({r["source"] for r in records}):
        source_records = [r for r in records if r["source"] == source]
        metrics = calc_metrics(source_records)
        metrics["n"] = len(source_records)
        labels = [r["label"] for r in source_records]
        if len(set(labels)) == 1:
            if labels[0] == 1:  # all fake
                metrics["detection_rate"] = metrics["accuracy"]
                if wiki_records:
                    pair_m = calc_metrics(wiki_records + source_records)
                    metrics["pairwise_auc"] = pair_m["auc_roc"]
                    metrics["pairwise_f1"] = pair_m["f1"]
            else:  # all real (wiki)
                metrics["false_alarm_rate"] = 1.0 - metrics["accuracy"]
        by_source[source] = metrics
    return by_source


# Real-vs-one-fake AUC/F1 per fake source - more interpretable than global AUC when class ratios vary
def pair_metrics(records, real_source="wiki"):
    fake_sources = sorted({r["source"] for r in records if r["source"] != real_source})
    pairwise = {}
    for fake_source in fake_sources:
        subset = [r for r in records if r["source"] in {real_source, fake_source}]
        if subset:
            pairwise[f"{real_source}_vs_{fake_source}"] = {
                "sources": [real_source, fake_source],
                "n": len(subset),
                **calc_metrics(subset),
            }
    return pairwise
