# %%
from dataclasses import dataclass

from matplotlib import pyplot as plt
from transformers import PreTrainedModel, PreTrainedTokenizerBase

from deception_detection.activations import Activations
from deception_detection.data import (
    AILiarDataset,
    AlpacaDataset,
    PairedDialogueDataset,
    RepeHonestDishonestFacts,
    WerewolfPaired,
)
from deception_detection.detectors import (
    MMSDetector,
)
from deception_detection.tokenized_data import TokenizedDataset, get_tokenized_dataset_pair
from deception_detection.utils import get_gemma_model_and_tokenizer


# %%
@dataclass
class PreparedDataset:
    dataset: PairedDialogueDataset
    honest_acts: Activations
    deceptive_acts: Activations

    def get_train_acts(self, val_fraction: float = 0.2):
        n_val = int(val_fraction * len(self.dataset))
        return self.honest_acts[:-n_val], self.deceptive_acts[:-n_val]

    def get_val_acts(self, val_fraction: float = 0.2):
        n_val = int(val_fraction * len(self.dataset))
        return self.honest_acts[-n_val:], self.deceptive_acts[-n_val:]


def prepare_dataset(
    dataset: PairedDialogueDataset,
    tokenizer: PreTrainedTokenizerBase,
    model: PreTrainedModel,
    batch_size: int = 16,
) -> PreparedDataset:
    honest_toks, deceptive_toks = get_tokenized_dataset_pair(
        dataset, tokenizer, enforce_equal_detect_lengths=True
    )

    honest_acts = Activations(model, honest_toks, batch_size=batch_size)
    deceptive_acts = Activations(model, deceptive_toks, batch_size=batch_size)

    return PreparedDataset(
        dataset=dataset,
        honest_acts=honest_acts,
        deceptive_acts=deceptive_acts,
    )


def get_alpaca_acts(tokenizer: PreTrainedTokenizerBase, model: PreTrainedModel) -> Activations:
    alpaca_dataset = AlpacaDataset()[:50]
    alpaca_toks = TokenizedDataset.from_dialogue_list(alpaca_dataset.dialogues, tokenizer)[:]
    alpaca_acts = Activations(model, alpaca_toks, batch_size=10)
    return alpaca_acts


# %%
model, tokenizer = get_gemma_model_and_tokenizer(low_mem=False)
alpaca_acts = get_alpaca_acts(tokenizer, model)

# %%
repe_dataset = prepare_dataset(RepeHonestDishonestFacts(), tokenizer, model)
ai_liar_dataset = prepare_dataset(AILiarDataset(detect_on_full_answer=True), tokenizer, model)
werewolf_dataset = prepare_dataset(WerewolfPaired(), tokenizer, model)


# %%
def get_detector():
    n_layers = 42
    detect_layers = list(range(int(0.25 * n_layers), int(0.75 * n_layers)))
    detector = MMSDetector(layers=detect_layers)
    honest_acts, deceptive_acts = repe_dataset.get_train_acts()
    detector.fit(deceptive_acts, honest_acts)
    return detector


detector = get_detector()

# %%
honest_scores_by_dialogue = detector.scores_by_dialogue(ai_liar_dataset.honest_acts)
deceptive_scores_by_dialogue = detector.scores_by_dialogue(ai_liar_dataset.deceptive_acts)


def get_range(*arrays: list[float]) -> tuple[float, float]:
    return (
        min(min(array) for array in arrays),
        max(max(array) for array in arrays),
    )


def plot_hist(honest_scores: list[float], deceptive_scores: list[float], ax: plt.Axes):
    histrange = get_range(honest_scores, deceptive_scores)
    ax.hist(honest_scores, bins=10, alpha=0.5, label="Honest", range=histrange)
    ax.hist(deceptive_scores, bins=10, alpha=0.5, label="Deceptive", range=histrange, color="r")


def max_mean_hists(dataset: PreparedDataset):
    honest_scores_by_dialogue = detector.scores_by_dialogue(dataset.honest_acts)
    deceptive_scores_by_dialogue = detector.scores_by_dialogue(dataset.deceptive_acts)

    _, axs = plt.subplots(2, 1, figsize=(6, 5), sharex=True)

    # max
    honest_max_scores = [score.max().item() for score in honest_scores_by_dialogue]
    deceptive_max_scores = [score.max().item() for score in deceptive_scores_by_dialogue]
    plot_hist(honest_max_scores, deceptive_max_scores, axs[0])  # type: ignore
    axs[0].set_title("Max score per dialogue")  # type: ignore

    # mean
    honest_mean_scores = [score.mean().item() for score in honest_scores_by_dialogue]
    deceptive_mean_scores = [score.mean().item() for score in deceptive_scores_by_dialogue]
    plot_hist(honest_mean_scores, deceptive_mean_scores, axs[1])  # type: ignore
    axs[1].set_title("Mean score per dialogue")  # type: ignore
    axs[1].set_xlabel("Score (→ Deceptive)")  # type: ignore
    axs[1].legend()  # type: ignore
    plt.suptitle(dataset.dataset.name)


max_mean_hists(ai_liar_dataset)
max_mean_hists(werewolf_dataset)


# %%
alpaca_threshold = detector.score(alpaca_acts).quantile(0.99).item()


def proportion_above_threshold_hist(dataset: PreparedDataset):
    def get_proportion(acts: Activations):
        scores = detector.scores_by_dialogue(acts)
        return [(score > alpaca_threshold).float().mean().item() for score in scores]

    honest_proportions = get_proportion(dataset.honest_acts)
    deceptive_proportions = get_proportion(dataset.deceptive_acts)
    plt.figure(figsize=(5, 3))
    plot_hist(honest_proportions, deceptive_proportions, plt.gca())
    plt.title(f"{dataset.dataset.name} - proportion above threshold")


proportion_above_threshold_hist(ai_liar_dataset)
proportion_above_threshold_hist(werewolf_dataset)


# %%
def plot_ecdfs(dataset: PreparedDataset):
    plt.figure(figsize=(6, 4))
    honest_scores_by_dialogue = detector.scores_by_dialogue(dataset.honest_acts)
    deceptive_scores_by_dialogue = detector.scores_by_dialogue(dataset.deceptive_acts)
    for score in deceptive_scores_by_dialogue:
        plt.ecdf(score, color="r", alpha=0.25)
    for score in honest_scores_by_dialogue:
        plt.ecdf(score, color="blue", alpha=0.25)
    plt.axvline(alpaca_threshold, color="black", linestyle="--")
    plt.title(dataset.dataset.name)
    plt.ecdf(detector.score(alpaca_acts), color="black", alpha=1)


plot_ecdfs(ai_liar_dataset)
plot_ecdfs(werewolf_dataset)
# %%
