# %%
from dataclasses import dataclass

import numpy as np
import pandas as pd
from transformers import PreTrainedModel, PreTrainedTokenizerBase

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.control import AlpacaDataset
from deception_detection.data.geometry_of_truth import Cities
from deception_detection.data.repe import RepeHonestDishonestFacts
from deception_detection.detectors import (
    DirectionDetector,
    LATDetector,
    LogisticRegressionDetector,
    MMSDetector,
)
from deception_detection.metrics import (
    get_auroc,
    get_tpr_at_fpr_from_control_dataset,
    get_tpr_at_fpr_from_paired_dataset,
)
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


# %%
@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 train_detector(
    prepared_dataset: PreparedDataset,
    detector_class: type[DirectionDetector],
    detector_kwargs: None | dict[str, list[int]] = None,
) -> DirectionDetector:
    if detector_kwargs is None:
        detector_kwargs = {"layers": [16]}
    detector = detector_class(**detector_kwargs)
    honest_acts, deceptive_acts = prepared_dataset.get_train_acts()
    detector.fit(deceptive_acts, honest_acts)
    return detector


def evaluate_detector(
    detector: DirectionDetector,
    prepared_dataset: PreparedDataset,
    fpr: float,
    alpaca_threshold: float,
    train_set: bool = False,
) -> tuple[float, float, float, float]:
    if train_set:
        honest_acts, deceptive_acts = prepared_dataset.get_val_acts()
    else:
        honest_acts, deceptive_acts = prepared_dataset.honest_acts, prepared_dataset.deceptive_acts
    honest_scores = detector.mean_scores_by_dialogue(honest_acts)
    deceptive_scores = detector.mean_scores_by_dialogue(deceptive_acts)
    auroc = get_auroc(honest_scores, deceptive_scores)
    tpr_neg_dataset = get_tpr_at_fpr_from_paired_dataset(honest_scores, deceptive_scores, fpr)
    tpr_alpaca = get_tpr_at_fpr_from_control_dataset(deceptive_scores, alpaca_threshold)
    fpr_alpaca_neg_dataset = get_tpr_at_fpr_from_control_dataset(honest_scores, alpaca_threshold)
    return auroc, tpr_neg_dataset, tpr_alpaca, fpr_alpaca_neg_dataset


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


def get_alpaca_threshold(detector: DirectionDetector, alpaca_acts: Activations, fpr: float):
    alpaca_scores = detector.mean_scores_by_dialogue(alpaca_acts)
    return float(np.percentile(alpaca_scores, 100 * (1 - fpr)))


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

# %%
datasets = {
    "RepeHonestDishonestFacts": RepeHonestDishonestFacts()[:200],
    "Cities": Cities()[:200],
    "AILiarWithAnswers": AILiarDataset(detect_on_full_answer=True),
    "InsiderTrading": PairedDialogueDataset.from_rollout_dialogues(
        load_rollout_dialogues("../../data/insider_trading/rollouts5.json")
    ),
}
prepared_datasets = {
    name: prepare_dataset(dataset, tokenizer, model) for name, dataset in datasets.items()
}

# %%
detector_types = [MMSDetector, LATDetector, LogisticRegressionDetector]
df = pd.DataFrame(
    columns=[
        "Detector",
        "Train Dataset",
        "Test Dataset",
        "AUROC",
        "TPR@1%FPR_neg",
        "TPR@1%FPR_alpaca",
        "FPR@1%FPR_alpaca",
    ]  # type: ignore
)
for detector_type in detector_types:
    print()
    print(f"==================== {detector_type.__name__} ====================")
    for train_name, train_set in prepared_datasets.items():
        detector = train_detector(train_set, detector_type, {"layers": list(range(10, 20))})
        alpaca_threshold = get_alpaca_threshold(detector, alpaca_acts, 0.01)
        for test_name, test_set in prepared_datasets.items():
            auroc, tpr_neg_dataset, tpr_alpaca, fpr_alpaca_neg_dataset = evaluate_detector(
                detector, test_set, fpr=0.01, alpaca_threshold=alpaca_threshold
            )
            row = [
                detector_type.__name__,
                train_name,
                test_name,
                auroc,
                tpr_neg_dataset,
                tpr_alpaca,
                fpr_alpaca_neg_dataset,
            ]
            df.loc[len(df)] = row
            print(
                "Train:",
                train_name,
                "Test:",
                test_name,
                "AUROC:",
                auroc,
                "TPR@1%FPR_neg:",
                tpr_neg_dataset,
                "TPR@1%FPR_alpaca:",
                tpr_alpaca,
                "FPR@1%FPR_alpaca:",
                fpr_alpaca_neg_dataset,
            )
            # print_tokens(detector, test_set, threshold)

df.to_csv("results.csv", index=False)

# %%

# def print_tokens(
#     detector: MMSDetector,
#     prepared_dataset: PreparedDataset,
#     threshold: float,
#     indices_to_show: list[int] = [0, 1, 2],
# ):
#     honest_scores_by_token = detector.scores_by_dialogue(prepared_dataset.honest_all_acts)
#     deceptive_scores_by_token = detector.scores_by_dialogue(prepared_dataset.deceptive_all_acts)

#     for idx in indices_to_show:
#         honest_str_tokens = prepared_dataset.honest_all_acts.tokenized_dataset.str_tokens[idx]
#         deceptive_str_tokens = prepared_dataset.deceptive_all_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][-honest_length:]
#         deceptive_score = deceptive_scores_by_token[idx][-deceptive_length:]
#         honest_to_plot = 2 * (honest_score > threshold) - 1
#         deceptive_to_plot = 2 * (deceptive_score > threshold) - 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",
#             )
#         )
# %%
for d in prepared_datasets.values():
    d.deceptive_acts.tokenized_dataset.display_detection_mask(tokenizer=tokenizer, idx=0)
    print()

# %%
