"""Analysis of how the relu threshold changes ROC AUC"""

import matplotlib.pyplot as plt
import torch

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.metrics import get_auroc
from deception_detection.prompt_scorer import ReluPromptScorer
from deception_detection.scores import Scores
from deception_detection.types import Label
from deception_detection.utils import REPO_ROOT


def get_aurocs_vs_honest(
    scores: Scores, scorer: ReluPromptScorer, thresholds: list[float]
) -> list[float]:
    """Calculate AUROC scores comparing honest vs deceptive responses."""
    aurocs = []
    for t in thresholds:
        scorer.relu_threshold = t
        aurocs.append(scorer.auroc(scores))
    return aurocs


def get_aurocs_vs_alpaca(
    scores: Scores, scorer: ReluPromptScorer, thresholds: list[float]
) -> list[float]:
    """Calculate AUROC scores comparing deceptive responses vs Alpaca baseline."""
    deceptive_scores = scores.get_scores_by_label()[Label.DECEPTIVE]
    aurocs = []
    for t in thresholds:
        scorer.relu_threshold = t
        # we don't use scorer.auroc_alpaca since it caches alpaca values
        control_honest = scorer.reduce(scorer.control_scores)
        dataset_deceptive = scorer.reduce(deceptive_scores.scores)
        aurocs.append(get_auroc(control_honest, dataset_deceptive))
    return aurocs


def main() -> None:
    """Main analysis function that loads experiment and creates plots."""
    # Load experiment
    cfg = ExperimentConfig.from_path(REPO_ROOT / "results/33/repe")
    experiment = Experiment(cfg)

    # Get Alpaca scores and initialize scorer
    # alpaca_scores = torch.cat(experiment.get_control_scores()["alpaca__plain"].scores)
    scorer = ReluPromptScorer(experiment, fprs=[0.05, 0.01, 0.002])
    init_threshold = scorer.relu_threshold

    # Setup thresholds and create figure with two subplots
    thresholds: list[float] = torch.linspace(-6, 6, 20).tolist()
    _, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 5), sharex=True)

    # Plot AUROCs vs honest
    for dataset_name, scores in experiment.get_eval_scores().items():
        if dataset_name == "repe_honesty__plain" or dataset_name.startswith("ai_audit"):
            continue
        aurocs = get_aurocs_vs_honest(scores, scorer, thresholds)
        ax1.plot(thresholds, aurocs, label=dataset_name)

    ax1.axvline(init_threshold, c="k", linestyle="--", label="default threshold")
    ax1.set_ylabel("AUROC vs honest")
    # ax1.set_ylim(0.95, 1)

    # Plot AUROCs vs alpaca
    for dataset_name, scores in experiment.get_eval_scores().items():
        if dataset_name == "repe_honesty__plain":
            continue
        aurocs = get_aurocs_vs_alpaca(scores, scorer, thresholds)
        ax2.plot(thresholds, aurocs, label=dataset_name.replace("__", "\n"))

    ax2.axvline(init_threshold, c="k", linestyle="--", label="default threshold")
    ax2.set_xlabel("ReLU threshold")
    ax2.set_ylabel("AUROC vs alpaca")
    ax2.legend()
    # ax2.set_ylim(0.99, 1)

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()
