# %%
from typing import Any

import matplotlib.pyplot as plt
import torch

from deception_detection.activations import Activations, SplitPairedActivations
from deception_detection.data import AILiarDataset, RepeHonestDishonestFacts
from deception_detection.sae import get_sae_acts
from deception_detection.utils import REPO_ROOT, get_model_and_tokenizer

OUT_FILE = REPO_ROOT / "deception_detection" / "notebooks" / "out" / "latent_selection"
OUT_FILE.mkdir(parents=True, exist_ok=True)
# %%

model, tokenizer = get_model_and_tokenizer("gemma", 9)

repe_acts = SplitPairedActivations.from_paired_dataset(
    dataset=RepeHonestDishonestFacts(),
    model=model,
    tokenizer=tokenizer,
    layers=[31],
    batch_size=5,
)
repe_sae_acts_honest = get_sae_acts(repe_acts.honest_acts)
repe_sae_acts_dishonest = get_sae_acts(repe_acts.deceptive_acts)


# %%
def frequency_latent_active(acts: Activations):
    latents_active = acts.get_masked_activations() > 0
    return torch.mean(latents_active.float(), dim=(0, 1))


repe_sae_on_freq_honest = frequency_latent_active(repe_sae_acts_honest)
repe_sae_on_freq_dishonest = frequency_latent_active(repe_sae_acts_dishonest)


# %%
def scatter_plot(honest_freq: torch.Tensor, deceptive_freq: torch.Tensor, **kwargs: Any):
    _ = plt.figure(figsize=(5, 5))
    kwargs = dict(alpha=0.1, s=2, c="k") | kwargs
    plt.scatter(honest_freq, deceptive_freq, **kwargs)
    plt.xlim(0, 1)
    plt.ylim(0, 1)
    plt.xlabel("Honest")
    plt.ylabel("Deceptive")
    plt.plot([0, 1], [0, 1], c="k", alpha=0.5, linewidth=0.5)


scatter_plot(repe_sae_on_freq_honest, repe_sae_on_freq_dishonest)
plt.savefig(OUT_FILE / "repe_all.png")

# %%
ail_acts = SplitPairedActivations.from_paired_dataset(
    dataset=AILiarDataset(variant="original_with_answers"),
    model=model,
    tokenizer=tokenizer,
    layers=[31],
    batch_size=5,
)
ail_sae_acts_honest = get_sae_acts(ail_acts.honest_acts)
ail_sae_acts_dishonest = get_sae_acts(ail_acts.deceptive_acts)

# %%
ail_sae_on_freq_honest = frequency_latent_active(ail_sae_acts_honest)
ail_sae_on_freq_dishonest = frequency_latent_active(ail_sae_acts_dishonest)

scatter_plot(ail_sae_on_freq_honest, ail_sae_on_freq_dishonest)
plt.savefig(OUT_FILE / "ail_all.png")


# %%
repe_deceptive_predictor_mask = (repe_sae_on_freq_dishonest - repe_sae_on_freq_honest) > 0.1
repe_honest_predictor_mask = (repe_sae_on_freq_honest - repe_sae_on_freq_dishonest) > 0.1

ail_deceptive_predictor_mask = (ail_sae_on_freq_dishonest - ail_sae_on_freq_honest) > 0.1
ail_honest_predictor_mask = (ail_sae_on_freq_honest - ail_sae_on_freq_dishonest) > 0.1


selected_latents_mask = (ail_deceptive_predictor_mask & repe_deceptive_predictor_mask) | (
    ail_honest_predictor_mask & repe_honest_predictor_mask
)
selected_latents: list[int] = selected_latents_mask.nonzero().squeeze().tolist()
print(selected_latents)


# %%
def scatter_plot_selected(
    honest: torch.Tensor,
    deceptive: torch.Tensor,
    selected_latents: list[int],
):
    _ = plt.figure(figsize=(5, 5))
    selected_honest = honest[selected_latents]
    selected_deceptive = deceptive[selected_latents]
    plt.scatter(selected_honest, selected_deceptive, s=4, c="k")
    plt.xlim(0, 1)
    plt.ylim(0, 1)
    plt.xlabel("Honest")
    plt.ylabel("Deceptive")
    plt.plot([0, 1], [0, 1], c="k", alpha=0.5, linewidth=0.5)
    for i, x, y in zip(selected_latents, selected_honest, selected_deceptive, strict=True):
        dx, dy = 0.02, 0.02
        plt.text(x.item() + dx, y.item() + dy, str(i), fontsize=8, ha="left", va="center")


scatter_plot_selected(repe_sae_on_freq_honest, repe_sae_on_freq_dishonest, selected_latents)
plt.title("repe_honesty")
plt.savefig(OUT_FILE / "repe_selected.png")

scatter_plot_selected(ail_sae_on_freq_honest, ail_sae_on_freq_dishonest, selected_latents)
plt.title("ai_liar_original_with_answers")
plt.savefig(OUT_FILE / "ail_selected.png")
# %%
