# %%

import matplotlib.pyplot as plt
import torch
from jaxtyping import Float
from torch import Tensor
from transformers import PreTrainedTokenizerBase

from deception_detection.activations import Activations
from deception_detection.data import CarSalesman
from deception_detection.notebooks.old.utils import get_repe_detector
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.utils import get_gemma_model_and_tokenizer

torch.set_grad_enabled(False)


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

# %%
car_data = CarSalesman()
car_toks = TokenizedDataset.from_dialogue_list(car_data.dialogues, tokenizer)
car_toks.display_detection_mask(0, tokenizer)

car_acts = Activations(model, car_toks, save_logits=True)
honesty_scores = detector.score(car_acts)


# %%
@torch.no_grad()
def get_logit_diff(
    acts: Activations,
    tokenizer: PreTrainedTokenizerBase,
    honest_completion: str,
    deceptive_completion: str,
) -> Float[Tensor, " batch"]:
    """Get the logit difference between honest and deceptive completion tokens.

    Extracts logits from activations and computes the difference between the logit of the honest
    and deceptive completion tokens. Also checks that the probabilities of the completions are
    reasonably high and thus the logit differences between these completions are representative.

    Args:
        acts: Activations object (used to get the masked logits)
        tokenizer: Tokenizer.
        honest_completion: Honest completion text (must be a single token).
        deceptive_completion: Deceptive completion text (must be a single token).

    Returns:
        Logit differences between honest and deceptive completions.
    """
    logits = acts.get_masked_logits()

    def encode_one(text: str) -> int:
        """Encode text consisting of a single token."""
        toks = tokenizer.encode(text, add_special_tokens=False)
        assert len(toks) == 1, f"Expected one token, got {len(toks)}"
        return toks[0]

    honest_tok_id = encode_one(honest_completion)
    deceptive_tok_id = encode_one(deceptive_completion)

    sum_probs = logits.softmax(dim=-1)[:, [honest_tok_id, deceptive_tok_id]].sum(dim=-1).cpu()
    assert (sum_probs > 1e-3).all(), f"Probabilities of completions are very low: {sum_probs}"

    logit_diffs = (logits[:, honest_tok_id] - logits[:, deceptive_tok_id]).cpu()
    return logit_diffs


logit_diffs = get_logit_diff(car_acts, tokenizer, "Yes", "No")

# %%

plt.gcf().set_size_inches(8, 6)
plt.scatter(logit_diffs, honesty_scores, c="black", s=8)
plt.xlabel("Logit difference (more honest →)")
plt.ylabel("Detector score (more honest →)")
for i, descriptor in enumerate(car_data.adjectives):
    plt.annotate(
        descriptor.split()[1],
        xy=(logit_diffs[i].item(), honesty_scores[i].item()),
        xytext=(5, -5),
        textcoords="offset points",
        fontsize=8,
    )
plt.xlim(None, 10.5)
