import torch 
from collections import defaultdict
from typing import cast

import numpy as np
from jaxtyping import Float
from sklearn.metrics import roc_auc_score, roc_curve, balanced_accuracy_score, accuracy_score
from torch import Tensor

from deception_detection.log import logger
from deception_detection.types import Label
"""
Note that our probes are trained such that "honest" is positive and "deceptive" is negative.    
But we want to measure the performance of the model in detecting deception.
Therefore, we need to flip the sign of the scores to get the correct AUROC (or alternatively,
we can flip the labels).
"""


def get_auroc_manual(
    honest_scores: Float[Tensor, " n_dialogues"],
    deceptive_scores: Float[Tensor, " n_dialogues"],
    N: int = 1000,
):
    """
    Stefan's implementation of auroc. Deprecated, but kept for reference.
    """
    eps = 1e-1
    small = min(honest_scores.min().item(), deceptive_scores.min().item()) - eps
    big = max(honest_scores.max().item(), deceptive_scores.max().item()) + eps
    thresholds = np.linspace(small, big, N)
    false_positives: dict[int, int] = defaultdict(int)
    true_positives: dict[int, int] = defaultdict(int)
    true_negatives: dict[int, int] = defaultdict(int)
    false_negatives: dict[int, int] = defaultdict(int)

    for honest_score, deceptive_score in zip(honest_scores, deceptive_scores, strict=True):
        # Higher score = more deceptive
        for i, threshold in enumerate(thresholds):
            # Honest prompt (negative)
            if honest_score <= threshold:
                # Detect honest (negative)
                true_negatives[i] += 1
            else:
                # Detect dishonest (positive)
                false_positives[i] += 1
            # Dishonest prompt (positive)
            if deceptive_score <= threshold:
                # Detect honest (negative)
                false_negatives[i] += 1
            else:
                # Detect dishonest (positive)
                true_positives[i] += 1

    tprs = [true_positives[i] / (true_positives[i] + false_negatives[i]) for i in range(N)]
    fprs = [false_positives[i] / (false_positives[i] + true_negatives[i]) for i in range(N)]

    # flip as fprs are in descending order
    auroc = -1 * np.trapz(tprs, fprs)

    return auroc


def get_true_and_scores(
    honest_scores: Float[Tensor, " n_dialogues"], deceptive_scores: Float[Tensor, " n_dialogues"]
):
    y_true = np.array([1] * len(deceptive_scores) + [0] * len(honest_scores))
    y_scores = np.concatenate([deceptive_scores, honest_scores])
    return y_true, y_scores


def get_auroc(
    honest_scores: Float[Tensor, " n_dialogues"],
    deceptive_scores: Float[Tensor, " n_dialogues"],
    weights: Float[Tensor, " n_dialogues"] | None = None,
) -> float:
    y_true, y_scores = get_true_and_scores(honest_scores, deceptive_scores)

    if np.isnan(y_scores).any():
        logger.warning("NaN scores found in AUROC calculation.")
        y_true = y_true[~np.isnan(y_scores)]
        y_scores = y_scores[~np.isnan(y_scores)]
    auroc = roc_auc_score(y_true, y_scores, sample_weight=weights)
    return cast(float, auroc)


def get_fpr_tpr(
    honest_scores: Float[Tensor, " n_dialogues"], deceptive_scores: Float[Tensor, " n_dialogues"]
) -> tuple[list[float], list[float], list[float]]:
    y_true, y_scores = get_true_and_scores(honest_scores, deceptive_scores)
    
    if np.isnan(y_scores).any():
        logger.warning("NaN scores found in FPR/TPR calculation.")
        y_true = y_true[~np.isnan(y_scores)]
        y_scores = y_scores[~np.isnan(y_scores)]
    
    fpr, tpr, thresholds = roc_curve(y_true, y_scores, drop_intermediate=False)
    fpr = fpr.tolist()
    tpr = tpr.tolist()
    thresholds = thresholds.tolist()
    return fpr, tpr, thresholds


def get_tpr_at_fpr_from_paired_dataset(
    honest_scores: Float[Tensor, " n_dialogues"],
    deceptive_scores: Float[Tensor, " n_dialogues"],
    fpr: float,
) -> float:
    fprs, tprs, _ = get_fpr_tpr(honest_scores, deceptive_scores)
    idx = np.where(np.array(fprs) > fpr)[0][0]
    tpr = tprs[idx]
    return cast(float, tpr)


def get_tpr_at_fpr_from_control_dataset(
    deceptive_scores: Float[Tensor, " n_dialogues"],
    threshold: float,
):
    return (deceptive_scores > threshold).float().mean().item()


def get_balanced_accuracy(
    scores: Float[Tensor, " n_dialogues"], 
    labels: list[Label], 
    threshold: float
) -> float:
    """
    Calculate balanced accuracy at a given threshold.
    """

    y_true = np.array([1 if label == Label.DECEPTIVE else 0 for label in labels])
    y_scores = scores.numpy()

    assert len(y_true) == len(y_scores), "Mismatch between labels and scores"
    
    if np.isnan(y_scores).any():
        logger.warning("NaN scores found in balanced accuracy calculation.")
        y_true = y_true[~np.isnan(y_scores)]
        y_scores = y_scores[~np.isnan(y_scores)]
    
    y_pred = (y_scores > threshold).astype(int)

    if len(np.unique(y_true)) == 1:
        # single-class dataset detected,therefore using accuracy instead of balanced accuracy
        logger.info("Single-class dataset detected, using accuracy instead of balanced accuracy")
        return accuracy_score(y_true, y_pred)
    else:
        balanced_acc = balanced_accuracy_score(y_true, y_pred)
        return balanced_acc