# %%
from collections import defaultdict

import matplotlib.pyplot as plt
import numpy as np
from torch import Tensor

from deception_detection.activations import Activations
from deception_detection.data import BlatantlyBadAdviceDataset, EasyGeography, PairedDialogueDataset
from deception_detection.notebooks.old.utils import get_repe_detector
from deception_detection.tokenized_data import get_tokenized_dataset_pair
from deception_detection.utils import get_gemma_model_and_tokenizer

# %%
detector_method = "mms"
model, tokenizer = get_gemma_model_and_tokenizer()
honesty_detector = get_repe_detector(model, tokenizer, method=detector_method)


# %%
def get_scores(dataset: PairedDialogueDataset) -> tuple[Tensor, Tensor]:
    honest_toks, deceptive_toks = get_tokenized_dataset_pair(dataset, tokenizer)
    honest_acts, deceptive_acts = (
        Activations(model, honest_toks),
        Activations(model, deceptive_toks),
    )
    honest_scores = honesty_detector.mean_scores_by_dialogue(honest_acts)
    deceptive_scores = honesty_detector.mean_scores_by_dialogue(deceptive_acts)
    return honest_scores, deceptive_scores


# %%
def plot_score_scatter(dataset: PairedDialogueDataset, ax: plt.Axes, title: str = ""):
    honest_scores, deceptive_scores = get_scores(dataset)

    ax.scatter(range(len(honest_scores)), honest_scores, color="C0", label="Honest")
    ax.scatter(range(len(deceptive_scores)), deceptive_scores, color="C1", label="deceptive")

    ax.set_xlabel("Example dialogue")
    ax.set_ylabel("Token-averaged honesty score")
    ax.legend()

    fraction_honest_higher_than_dishonest = (honest_scores > deceptive_scores).float().mean()
    ax.set_title(f"{fraction_honest_higher_than_dishonest:.1%} intra-example accuracy", loc="right")


def generate_curve(dataset: PairedDialogueDataset, ax: plt.Axes, title: str = "", N: int = 3000):
    ax.set_xlabel("False positive rate")
    ax.set_ylabel("True positive rate")
    thresholds = np.linspace(-3, 3, 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)

    honest_scores, deceptive_scores = get_scores(dataset)
    for honest_score, deceptive_score in zip(honest_scores, deceptive_scores, strict=True):
        # Higher score = more honest
        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)]
    ax.plot(fprs, tprs, color="k")

    auroc = np.trapz(tprs, fprs)
    ax.set_title(f"AUROC={auroc:.2f}            {title}", loc="left")


#  %% Plot scores for an example

# plot_detector_scores(data[0]["honest"], tokenizer, honesty_detector)
# plot_detector_scores(data[0]["deceptive"], tokenizer, honesty_detector)

# %% Run benchmarks
fig, axes = plt.subplots(
    nrows=5, ncols=2, figsize=(9, 15), width_ratios=[1, 1.5], constrained_layout=True
)
assert isinstance(axes, np.ndarray)
fig.suptitle(f"{detector_method} evaluation")

geo_dataset = EasyGeography()
generate_curve(geo_dataset, axes[0, 0], title="easy geography")
plot_score_scatter(geo_dataset, axes[0, 1], title="easy geography")

advice_dataset = BlatantlyBadAdviceDataset()
generate_curve(advice_dataset, axes[1, 0], title="blatantly bad advice")
plot_score_scatter(advice_dataset, axes[1, 1], title="blatantly bad advice")

# %%

# TODO: impliment the dataset below

# with open(REPO_ROOT / "data/good_actions_bad_actions.json") as file:
#     good_data = []
#     bad_data = []
#     all_data = []
#     for d in json.load(file):
#         good_data.append(d["good"])
#         all_data.append(d["good"])
#         bad_data.append(d["bad"])
#         all_data.append(d["bad"])

#     generate_curve(all_data, axes[2, 0], title="good and bad actions")
#     plot_score_scatter(all_data, axes[2, 1], title="good and bad actions")

#     generate_curve(bad_data, axes[3, 0], title="bad actions")
#     plot_score_scatter(bad_data, axes[3, 1], title="bad actions")

#     generate_curve(good_data, axes[4, 0], title="good actions")
#     plot_score_scatter(good_data, axes[4, 1], title="good actions")


# %%
