"""
What I imagine as the interface via which you should be able to train and evaluate a deception
detection method.
"""

# %%
from typing import Any

import matplotlib.pyplot as plt

from deception_detection.activations import Activations
from deception_detection.data import AILiarDataset
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

# %%
model, tokenizer = get_gemma_model_and_tokenizer()
detector = get_repe_detector(model, tokenizer, "mms")

# %%
data = AILiarDataset()
honest_toks, deceptive_toks = get_tokenized_dataset_pair(data, tokenizer)
honest_toks.display_detection_mask(0, tokenizer)
honest_acts, deceptive_acts = Activations(model, honest_toks), Activations(model, deceptive_toks)

# %%
honest_scores = detector.score(honest_acts)
deceptive_scores = detector.score(deceptive_acts)


# %%
def scatter_with_diagonal(x: Any, y: Any, **kwargs: Any) -> None:
    """Square scatter plot with diagonal line"""
    plt.figure(figsize=(4, 4))
    plt.scatter(x, y, **kwargs)
    xmin, xmax = plt.xlim()
    ymin, ymax = plt.ylim()
    lims = [min(xmin, ymin), max(xmax, ymax)]
    plt.plot(lims, lims, color="black", linestyle="--")
    plt.xlim(lims)
    plt.ylim(lims)
    plt.gca().set_aspect("equal", adjustable="box")


scatter_with_diagonal(honest_scores, deceptive_scores, c="black", s=8)
plt.xlabel("Honesty score on honest completions")
plt.ylabel("Honesty score on deceptive completions")

# %%
