# # %%
# from pathlib import Path
# from typing import cast

# import matplotlib.pyplot as plt
# import torch

# from deception_detection.activations import Activations
# from deception_detection.experiment import Experiment, ExperimentConfig
# from deception_detection.rollouts import load_rollout_dialogues
# from deception_detection.tokenized_data import TokenizedDataset
# from deception_detection.types import Dialogue, Message, RolloutDialogue
# from deception_detection.utils import REPO_ROOT, get_gemma_model_and_tokenizer

# # %%
# folder = REPO_ROOT / "deception_detection/writeup/results/mms_repe_9b_20240816_094155"
# cfg = ExperimentConfig.from_path(Path(folder))
# model, tokenizer = get_gemma_model_and_tokenizer(cfg.size_model)
# experiment = Experiment(cfg, model=model, tokenizer=tokenizer)
# assert experiment.cfg.detect_only_start_of_turn

# # %%


# def maybe_append_assistant(dialogue: Dialogue) -> Dialogue:
#     if dialogue[-1].role != "assistant":
#         dialogue.append(Message("assistant", "", True))
#     return dialogue


# def compute_scores(rollouts: list[RolloutDialogue], experiment: Experiment) -> list[float]:
#     dialogues = [maybe_append_assistant(rollout.input_messages) for rollout in rollouts]
#     assert all(dialogue[-1].role == "assistant" for dialogue in dialogues)

#     tokenized_dataset = TokenizedDataset.from_dialogue_list(
#         dialogues, experiment.tokenizer, detect_k_prior_tokens=4, detect_only_start_of_turn=True
#     )
#     assert experiment.model is not None
#     acts = Activations(experiment.model, tokenized_dataset)
#     return [row_scores.mean().item() for row_scores in experiment.get_detector().score(acts).
# scores]


# data_dir = REPO_ROOT / "data"

# rollouts = {
#     "ai_liar_rollouts": load_rollout_dialogues(
#         data_dir / "how_to_catch_an_ai_liar/rollouts5_9b.json"
#     ),
#     "werewolf_rollouts": load_rollout_dialogues(data_dir / "werewolf/rollouts10_9b.json"),
#     "insider_trading": load_rollout_dialogues(data_dir / "insider_trading/rollouts5.json"),
# }
# for werewolf_prompt, role in zip(
#     rollouts["werewolf_rollouts"], WerewolfOpenResponse().roles, strict=False
# ):
#     werewolf_prompt.expect_honesty = role not in ["Werewolf", "Minion"]

# scores = {key: compute_scores(rollouts[key], experiment) for key in rollouts}

# # %%
# fig, axs = plt.subplots(1, 3, figsize=(7, 4), sharey=True, sharex=True)

# for ax, dataset in zip(
#     axs, ["ai_liar_rollouts", "werewolf_rollouts", "insider_trading"], strict=False
# ):
#     if dataset == "insider_trading":
#         c = "k"
#     else:
#         c = ["cornflowerblue" if r.expect_honesty else "firebrick" for r in rollouts[dataset]]

#     x = [1 - cast(float, r.honesty_percentage) for r in rollouts[dataset]]
#     # apply jitter
#     x = [xi + 0.02 * (torch.rand(1) - 0.5) for xi in x]
#     ax.scatter(x=x, y=scores[dataset], c=c, s=15, alpha=0.5)

#     ax.set_title(dataset)
#     ax.set_box_aspect(1)
#     ax.set_xlabel("Deception rate")

# axs[0].set_ylabel("Detector score")
# plt.tight_layout()


# # %%
