# %%
from typing import cast

import numpy as np
from circuitsvis.tokens import colored_tokens
from IPython.display import display

from deception_detection.activations import Activations
from deception_detection.data.ai_liar import AILiarDataset
from deception_detection.data.base import PairedDialogueDataset
from deception_detection.data.geometry_of_truth import Cities
from deception_detection.data.repe import RepeHonestDishonestFacts
from deception_detection.detectors import DirectionDetector, LogisticRegressionDetector
from deception_detection.metrics import get_auroc, get_fpr_tpr
from deception_detection.rollouts import load_rollout_dialogues
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()

# %% train a detector on repehonestdishonestfacts

dataset = RepeHonestDishonestFacts()
val_fraction = 0.2
n_val = int(val_fraction * len(dataset))
train_dataset = dataset[:-n_val]
val_dataset = train_dataset[-n_val:]

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 = LogisticRegressionDetector(layers=[16])
detector.fit(deceptive_train_acts, honest_train_acts)


# %%
def print_metrics(
    honest_acts: Activations, deceptive_acts: Activations, detector: DirectionDetector
):
    honest_scores = detector.max_scores_by_dialogue(honest_acts)
    deceptive_scores = detector.max_scores_by_dialogue(deceptive_acts)
    auroc = get_auroc(honest_scores, deceptive_scores)
    fprs, tprs, thresholds = get_fpr_tpr(honest_scores, deceptive_scores)
    idx = np.where(np.array(fprs) > 0.01)[0][0]
    tpr = tprs[idx]
    threshold = thresholds[idx]
    print(f"Validation AUROC: {auroc}")
    print(f"Validation TPR at 1% FPR: {tpr}")
    print(f"Threshold for 1% FPR: {threshold}")
    return cast(float, threshold)


def print_tokens(
    honest_acts: Activations,
    deceptive_acts: Activations,
    threshold: float,
    indices_to_show: list[int] | None = None,
    act_fn: bool = True,
):
    """
    TODO: move to a non-notebook filea
    """
    if indices_to_show is None:
        indices_to_show = [0, 1, 2]
    honest_scores_by_token = detector.scores_by_dialogue(honest_acts)
    deceptive_scores_by_token = detector.scores_by_dialogue(deceptive_acts)
    for idx in indices_to_show:
        honest_str_tokens = honest_acts.tokenized_dataset.str_tokens[idx]
        deceptive_str_tokens = deceptive_acts.tokenized_dataset.str_tokens[idx]
        honest_length = len(honest_str_tokens)
        deceptive_length = len(deceptive_str_tokens)
        honest_score = honest_scores_by_token[idx]
        deceptive_score = deceptive_scores_by_token[idx]
        honest_score = honest_score[-honest_length:]
        deceptive_score = deceptive_score[-deceptive_length:]
        if act_fn:
            honest_to_plot = 2 * (honest_score > threshold) - 1
            deceptive_to_plot = 2 * (deceptive_score > threshold) - 1
        else:
            # just normalise within the prompt
            small = min(honest_score.min(), deceptive_score.min())  # type: ignore
            big = max(honest_score.max(), deceptive_score.max())  # type: ignore
            honest_to_plot = 2 * ((honest_score - small) / (big - small)) - 1
            deceptive_to_plot = 2 * ((deceptive_score - small) / (big - small)) - 1
        print("Honest:")
        display(
            colored_tokens(
                honest_str_tokens,
                honest_to_plot,
                min_value=-1,
                max_value=1,
                positive_color="red",
                negative_color="green",
            )
        )
        print("Deceptive:")
        display(
            colored_tokens(
                deceptive_str_tokens,
                deceptive_to_plot,
                min_value=-1,
                max_value=1,
                positive_color="red",
                negative_color="green",
            )
        )


# %% evaluate the detector in distribution
honest_val_toks, deceptive_val_toks = get_tokenized_dataset_pair(
    val_dataset, tokenizer, enforce_equal_detect_lengths=True
)

honest_val_acts = Activations(model, honest_val_toks)
deceptive_val_acts = Activations(model, deceptive_val_toks)

print("In distribution evaluation:")
threshold = print_metrics(honest_val_acts, deceptive_val_acts, detector)

all_honest_val_toks, all_deceptive_val_toks = get_tokenized_dataset_pair(
    val_dataset, tokenizer, enforce_equal_detect_lengths=False, detect_all=True
)
honest_val_acts = Activations(model, all_honest_val_toks)
deceptive_val_acts = Activations(model, all_deceptive_val_toks)
print_tokens(honest_val_acts, deceptive_val_acts, threshold, act_fn=False)
print_tokens(honest_val_acts, deceptive_val_acts, threshold, act_fn=True)

# %% evaluate the detector on cities
cities_dataset = Cities()
honest_cities_toks, deceptive_cities_toks = get_tokenized_dataset_pair(
    cities_dataset, tokenizer, enforce_equal_detect_lengths=False
)
honest_cities_acts = Activations(model, honest_cities_toks)
deceptive_cities_acts = Activations(model, deceptive_cities_toks)
print("Out of distribution evaluation on Cities:")
threshold = print_metrics(honest_cities_acts, deceptive_cities_acts, detector)

all_honest_cities_toks, all_deceptive_cities_toks = get_tokenized_dataset_pair(
    cities_dataset, tokenizer, enforce_equal_detect_lengths=False, detect_all=True
)
honest_cities_acts = Activations(model, all_honest_cities_toks)
deceptive_cities_acts = Activations(model, all_deceptive_cities_toks)

# %%
print_tokens(honest_cities_acts, deceptive_cities_acts, threshold, act_fn=False)
print_tokens(honest_cities_acts, deceptive_cities_acts, threshold, act_fn=True)


# %% evaluate the detector on AILiar (non rollouts)
ai_liar_dataset = AILiarDataset()

honest_ai_liar_toks, deceptive_ai_liar_toks = get_tokenized_dataset_pair(
    ai_liar_dataset, tokenizer, enforce_equal_detect_lengths=False
)
honest_ai_liar_acts = Activations(model, honest_ai_liar_toks)
deceptive_ai_liar_acts = Activations(model, deceptive_ai_liar_toks)

print("Out of distribution evaluation on AILiar:")
threshold = print_metrics(honest_ai_liar_acts, deceptive_ai_liar_acts, detector)

all_honest_ai_liar_toks, all_deceptive_ai_liar_toks = get_tokenized_dataset_pair(
    ai_liar_dataset, tokenizer, enforce_equal_detect_lengths=False, detect_all=True
)
honest_ai_liar_acts = Activations(model, all_honest_ai_liar_toks)
deceptive_ai_liar_acts = Activations(model, all_deceptive_ai_liar_toks)
print_tokens(honest_ai_liar_acts, deceptive_ai_liar_acts, threshold)

# %% Evaluate the detector on AILiar with the full written answer

ai_liar_dataset = AILiarDataset(detect_on_full_answer=True)

honest_ai_liar_toks, deceptive_ai_liar_toks = get_tokenized_dataset_pair(
    ai_liar_dataset, tokenizer, enforce_equal_detect_lengths=False
)
honest_ai_liar_acts = Activations(model, honest_ai_liar_toks)
deceptive_ai_liar_acts = Activations(model, deceptive_ai_liar_toks)

print("Out of distribution evaluation on AILiar:")
threshold = print_metrics(honest_ai_liar_acts, deceptive_ai_liar_acts, detector)

all_honest_ai_liar_toks, all_deceptive_ai_liar_toks = get_tokenized_dataset_pair(
    ai_liar_dataset, tokenizer, enforce_equal_detect_lengths=False, detect_all=True
)
honest_ai_liar_acts = Activations(model, all_honest_ai_liar_toks)
deceptive_ai_liar_acts = Activations(model, all_deceptive_ai_liar_toks)
print_tokens(honest_ai_liar_acts, deceptive_ai_liar_acts, threshold)

# %% evalute on insider trading
rollout_dialogues = load_rollout_dialogues("../../data/insider_trading/rollouts5.json")
insider_trading_dataset = PairedDialogueDataset.from_rollout_dialogues(
    rollout_dialogues, name="insider_trading"
)
print(len(insider_trading_dataset))
honest_insider_trading_toks, deceptive_insider_trading_toks = get_tokenized_dataset_pair(
    insider_trading_dataset, tokenizer, enforce_equal_detect_lengths=False
)

honest_insider_trading_acts = Activations(model, honest_insider_trading_toks, batch_size=2)
deceptive_insider_trading_acts = Activations(model, deceptive_insider_trading_toks, batch_size=2)

print("Out of distribution evaluation on Insider Trading:")
threshold = print_metrics(honest_insider_trading_acts, deceptive_insider_trading_acts, detector)

all_honest_insider_trading_toks, all_deceptive_insider_trading_toks = get_tokenized_dataset_pair(
    insider_trading_dataset, tokenizer, enforce_equal_detect_lengths=False, detect_all=True
)
honest_insider_trading_acts = Activations(model, all_honest_insider_trading_toks, batch_size=2)
deceptive_insider_trading_acts = Activations(
    model, all_deceptive_insider_trading_toks, batch_size=2
)
print_tokens(honest_insider_trading_acts, deceptive_insider_trading_acts, threshold)


# %%
print_tokens(honest_insider_trading_acts, deceptive_insider_trading_acts, -0.5)
# %%
