"""
Train a probe on other datasets and test generalization to insider trading. This notebook is a mess.
"""

# %%
import plotly.express as px
import plotly.graph_objects as go
from circuitsvis.tokens import colored_tokens
from IPython.display import display

from deception_detection.activations import Activations
from deception_detection.data.base import RolloutDialogueDataset
from deception_detection.data.geometry_of_truth import Cities
from deception_detection.data.repe import RepeHonestDishonestFacts
from deception_detection.detectors import MMSDetector
from deception_detection.rollouts import load_rollout_dialogues
from deception_detection.tokenized_data import TokenizedDataset, get_tokenized_dataset_pair
from deception_detection.utils import get_gemma_model_and_tokenizer

model, tokenizer = get_gemma_model_and_tokenizer()

# %% train a detector
train_dataset = RepeHonestDishonestFacts()
honest_train_toks, deceptive_train_toks = get_tokenized_dataset_pair(
    train_dataset, tokenizer, enforce_equal_detect_lengths=True
)
honest_train_acts = Activations(model, honest_train_toks)
deceptive_train_acts = Activations(model, deceptive_train_toks)
detector = MMSDetector(layers=list(range(10, 20)))
detector.fit(honest_train_acts, deceptive_train_acts)

# %%
honest_train_toks.display_detection_mask(8, tokenizer)

# %% some other paired dataset
test_dataset = Cities()
honest_test_toks, deceptive_test_toks = get_tokenized_dataset_pair(
    test_dataset, tokenizer, detect_all=True
)
honest_test_acts = Activations(model, honest_test_toks)
deceptive_test_acts = Activations(model, deceptive_test_toks)
honest_scores = detector.scores_by_dialogue(honest_test_acts)
deceptive_scores = detector.scores_by_dialogue(deceptive_test_acts)
# %%
# print a sample honest and harmful completion
idx = 30
honest_str_tokens = honest_test_toks.str_tokens[idx][4:]
deceptive_str_tokens = deceptive_test_toks.str_tokens[idx][4:]
honest_length = len(honest_str_tokens)
deceptive_length = len(deceptive_str_tokens)
honest_score = honest_scores[idx][4:]
deceptive_score = deceptive_scores[idx][4:]
x_labels = []
for i, honest_label, dishonest_label in zip(
    list(range(len(honest_str_tokens))), honest_str_tokens, deceptive_str_tokens, strict=False
):
    if honest_label == dishonest_label:
        x_labels.append(f"{i}: {honest_label}")
    else:
        x_labels.append(f"{i}: {honest_label}(+)  / {dishonest_label}(-)")


fig = go.Figure()

fig.add_trace(
    go.Scatter(x=x_labels, y=honest_score[-honest_length:], mode="lines+markers", name="Honest")
)

fig.add_trace(
    go.Scatter(
        x=x_labels, y=deceptive_score[-deceptive_length:], mode="lines+markers", name="Deceptive"
    )
)

fig.update_layout(
    title="Honest vs Deceptive Probe Scores",
    xaxis_title="Tokens",
    yaxis_title="Score",
    legend_title="Type",
    xaxis_tickangle=-45,
)

# Show the plot
fig.show()

honest_score = honest_score[-honest_length:]
deceptive_score = deceptive_score[-deceptive_length:]

min_score = min(min(honest_score), min(deceptive_score))  # type: ignore
max_score = max(max(honest_score), max(deceptive_score))  # type: ignore

# normalise to [-1, 1]
honest_score = (honest_score - min_score) / (max_score - min_score) * 2 - 1
deceptive_score = (deceptive_score - min_score) / (max_score - min_score) * 2 - 1


print("Honest:")
display(
    colored_tokens(
        honest_str_tokens,
        honest_score,  # type: ignore
        min_value=-1,  # type: ignore
        max_value=1,  # type: ignore
        positive_color="green",
        negative_color="red",
    )  # type: ignore
)
print("Deceptive:")
display(
    colored_tokens(
        deceptive_str_tokens,
        deceptive_score,  # type: ignore
        min_value=-1,  # type: ignore
        max_value=1,  # type: ignore
        positive_color="green",
        negative_color="red",
    )
)


# %% Test on AI Liar rollouts
ai_liar_rollout_dialogues = load_rollout_dialogues("data/how_to_catch_an_ai_liar/rollouts.json")
ai_liar = RolloutDialogueDataset.from_rollout_dialogues(ai_liar_rollout_dialogues)
for dialogue in ai_liar.dialogues:
    for message in dialogue:
        message.detect = True
    # need to do this for some reason
    dialogue[-1].content = dialogue[-1].content.rstrip("\n")
    dialogue[-1].content = dialogue[-1].content.rstrip(" ")
tokenized_ai_liar = TokenizedDataset.from_dialogue_list(
    ai_liar.dialogues, tokenizer, detect_all=True
)
activations = Activations(model, tokenized_ai_liar, batch_size=2)

# %%
scores = detector.scores_by_dialogue(activations)
# %%
idx = 52
honesty = ai_liar.model_graded_honesty[idx]
str_tokens = tokenized_ai_liar.str_tokens[idx]
length = len(str_tokens)
trim_str_tokens = str_tokens[4:]
score = scores[idx][-length + 4 :]
norm_score = (score - min_score) / (max_score - min_score) * 2 - 1
px.line(norm_score).show()  # type: ignore
print(honesty)
colored_tokens(
    trim_str_tokens,
    norm_score,  # type: ignore
    min_value=-1,
    max_value=1,
    positive_color="green",
    negative_color="red",
)


# %% test dataset insider trading
insider_trading_rollout_dialogues = load_rollout_dialogues("data/insider_trading/rollouts.json")
insider_trading = RolloutDialogueDataset.from_rollout_dialogues(insider_trading_rollout_dialogues)
for dialogue in insider_trading.dialogues:
    for message in dialogue:
        message.detect = True
    # need to do this for some reason
    dialogue[-1].content = dialogue[-1].content.rstrip("\n")
    dialogue[-1].content = dialogue[-1].content.rstrip(" ")
tokenized_insider_trading = TokenizedDataset.from_dialogue_list(
    insider_trading.dialogues, tokenizer, detect_all=True
)
activations = Activations(model, tokenized_insider_trading, batch_size=2)

# %%
scores = detector.scores_by_dialogue(activations)
print(
    list(
        zip(
            insider_trading.model_graded_honesty,
            list(range(len(insider_trading.model_graded_honesty))),
            strict=False,
        )
    )  # type: ignore
)  # type: ignore

# %%
idx = 14
honesty = insider_trading.model_graded_honesty[idx]
str_tokens = tokenized_insider_trading.str_tokens[idx][4:]
length = len(str_tokens)
score = scores[idx][-length + 4 :]
norm_score = (score - min_score) / (max_score - min_score) * 2 - 1
px.line(norm_score).show()  #    type: ignore
print(honesty)
colored_tokens(
    str_tokens,
    norm_score,  # type: ignore
    min_value=-1,
    max_value=1,
    positive_color="green",
    negative_color="red",  # type: ignore
)


# %%
