# %%
from collections.abc import Sequence
from dataclasses import dataclass
from itertools import product

import pandas as pd
import plotly.graph_objects as go
import torch
from tqdm import tqdm
from transformers import PreTrainedModel, PreTrainedTokenizer

from deception_detection.activations import Activations
from deception_detection.data import (
    AILiarDataset,
    Animals,
    Cities,
    Elements,
    Facts,
    Inventions,
    NegCities,
    NegSpEnTrans,
    PairedDialogueDataset,
    PromptedSandbaggingDataset,
    RepeHonestDishonestFacts,
    SpEnTrans,
    SycophancyNLP,
    SycophancyPhilPapers,
    SycophancyPoliticalTypology,
    TruthfulQA,
    WerewolfPaired,
)
from deception_detection.data.base import RolloutDialogueDataset
from deception_detection.detectors import MMSDetector
from deception_detection.metrics import get_auroc
from deception_detection.results import (
    GeneralizationResult,
    load_generalization_results,
    save_generalization_results,
)
from deception_detection.rollouts import load_rollout_dialogues
from deception_detection.tokenized_data import TokenizedDataset, get_tokenized_dataset_pair
from deception_detection.types import Dialogue
from deception_detection.utils import get_gemma_model_and_tokenizer


@dataclass
class PreparedDataset:
    name: str
    train: tuple[Activations, Activations]
    test: tuple[Activations, Activations]


def prepare_rollout_datasets(
    datasets: Sequence[RolloutDialogueDataset],
    tokenizer: PreTrainedTokenizer,
    model: PreTrainedModel,
):
    """
    Prepare datasets for training and testing.

    Args:
    datasets (List): List of dataset objects to prepare.
    tokenizer: Tokenizer object to use for tokenization.
    model: Model object to use for generating activations.

    Returns:
    List[Dict]: List of dictionaries containing prepared datasets.
    """
    prepared_datasets: list[PreparedDataset] = []

    for dataset in tqdm(datasets, desc="Preparing datasets"):
        honest_dialogues: list[Dialogue] = []
        deceptive_dialogues: list[Dialogue] = []
        for honest, dialogue in zip(dataset.model_graded_honesty, dataset.dialogues, strict=False):
            if honest == "honest":
                honest_dialogues.append(dialogue)
            elif honest == "deceptive":
                deceptive_dialogues.append(dialogue)
        honest_toks = TokenizedDataset.from_dialogue_list(honest_dialogues, tokenizer)
        deceptive_toks = TokenizedDataset.from_dialogue_list(deceptive_dialogues, tokenizer)
        honest_acts = Activations(model, honest_toks, batch_size=2)
        deceptive_acts = Activations(model, deceptive_toks, batch_size=2)
        prepared_datasets.append(
            PreparedDataset(
                name=dataset.name,
                train=(honest_acts, deceptive_acts),
                test=(honest_acts, deceptive_acts),
            )
        )
    return prepared_datasets


def prepare_datasets(
    datasets: Sequence[PairedDialogueDataset],
    tokenizer: PreTrainedTokenizer,
    model: PreTrainedModel,
    train_ratio: float = 0.7,
) -> list[PreparedDataset]:
    """
    Prepare datasets for training and testing.

    Args:
    datasets (List): List of dataset objects to prepare.
    tokenizer: Tokenizer object to use for tokenization.
    model: Model object to use for generating activations.
    train_ratio (float): Ratio of data to use for training (default: 0.7).

    Returns:
    List[Dict]: List of dictionaries containing prepared datasets.
    """
    prepared_datasets: list[PreparedDataset] = []

    for dataset in tqdm(datasets, desc="Preparing datasets"):
        n_train = int(len(dataset) * train_ratio)
        train_data, test_data = dataset[:n_train], dataset[n_train:]

        honest_train_toks, deceptive_train_toks = get_tokenized_dataset_pair(
            train_data, tokenizer, enforce_equal_detect_lengths=True
        )
        honest_test_toks, deceptive_test_toks = get_tokenized_dataset_pair(
            test_data, tokenizer, enforce_equal_detect_lengths=True
        )

        honest_train_acts = Activations(model, honest_train_toks)
        deceptive_train_acts = Activations(model, deceptive_train_toks)
        honest_test_acts = Activations(model, honest_test_toks)
        deceptive_test_acts = Activations(model, deceptive_test_toks)

        prepared_datasets.append(
            PreparedDataset(
                name=dataset.name,
                train=(honest_train_acts, deceptive_train_acts),
                test=(honest_test_acts, deceptive_test_acts),
            )
        )

    return prepared_datasets


def evaluate_datasets(
    train_datasets: list[PreparedDataset],
    test_datasets: list[PreparedDataset],
    layer_min: int = 10,
    layer_max: int = 20,
    method: str = "mms",
    flip: bool = False,
) -> list[GeneralizationResult]:
    """
    Train and evaluate MMSDetector on all pairs of datasets.
    """

    results: list[GeneralizationResult] = []

    for train_dataset, test_dataset in tqdm(
        product(train_datasets, test_datasets),
        desc="Evaluating dataset pairs",
    ):
        if method == "mms":
            detector = MMSDetector(layers=list(range(layer_min, layer_max)))
        else:
            raise NotImplementedError(f"Method {method} not implemented.")

        detector.fit(*train_dataset.train)

        if flip:
            assert detector.directions is not None
            detector.directions = -detector.directions

        honest_test_acts, deceptive_test_acts = test_dataset.test
        test_honest_scores = detector.scores_by_dialogue(honest_test_acts)
        test_honest_scores = torch.tensor([score.mean() for score in test_honest_scores])
        test_dishonest_scores = detector.scores_by_dialogue(deceptive_test_acts)
        test_dishonest_scores = torch.tensor([score.mean() for score in test_dishonest_scores])
        auroc = get_auroc(test_honest_scores, test_dishonest_scores)

        result = GeneralizationResult(
            train_dataset=train_dataset.name,
            test_dataset=test_dataset.name,
            model=method,
            layers=list(range(layer_min, layer_max)),
            auroc=auroc,
        )

        results.append(result)
    return results


def create_plot(
    results: list[GeneralizationResult], attribute: str, ordering: list[str]
) -> go.Figure:
    """
    Create a heatmap visualization of accuracy results.
    """

    # Create a DataFrame from the results
    df = pd.DataFrame(
        [
            {
                "train_dataset": r.train_dataset,
                "test_dataset": r.test_dataset,
                "model": r.model,
                "auroc": r.auroc,
            }
            for r in results
        ]
    )

    # Pivot the DataFrame to create the grid
    pivot_df = df.pivot(index="train_dataset", columns="test_dataset", values=attribute)

    # Reorder the columns and rows
    rows = pivot_df.index
    cols = pivot_df.columns

    row_ordering = [x for x in ordering if x in rows]
    col_ordering = [x for x in ordering if x in cols]

    pivot_df = pivot_df.loc[row_ordering, col_ordering]

    # Create the heatmap
    fig = go.Figure(
        data=go.Heatmap(
            z=pivot_df.values,  # type: ignore
            x=pivot_df.columns,  # type: ignore
            y=pivot_df.index,  # type: ignore
            colorscale="RdBU",
            zmin=0,
            zmax=1,
            text=pivot_df.values.round(2),  # type: ignore
            hoverinfo="text",
            texttemplate="%{text:.2f}",
        )
    )

    # Update the layout
    fig.update_layout(
        title="Generalization Results: AUROC Scores",
        xaxis_title="Test Dataset",
        yaxis_title="Train Dataset",
        width=1000,
        height=1000,
    )
    return fig


# %% load model
model, tokenizer = get_gemma_model_and_tokenizer()

insider_trading_dialogues = load_rollout_dialogues("data/insider_trading/rollouts.json")
insider_trading_dataset = RolloutDialogueDataset.from_rollout_dialogues(
    insider_trading_dialogues, name="Insider Trading"
)
for dialogue in insider_trading_dataset.dialogues:
    dialogue[-1].content = dialogue[-1].content.rstrip("\n")
    dialogue[-1].content = dialogue[-1].content.rstrip(" ")
prepared_insider_trading = prepare_rollout_datasets([insider_trading_dataset], tokenizer, model)  # type: ignore

datasets = [
    # Facts
    RepeHonestDishonestFacts(),
    Cities(),
    NegCities(),
    SpEnTrans(),
    NegSpEnTrans(),
    Animals(),
    Elements(),
    Facts(),
    Inventions(),
    # Imitative Falsehoods
    TruthfulQA(),
    # Sycophancy
    SycophancyNLP(),
    SycophancyPhilPapers(),
    SycophancyPoliticalTypology(),
    # Deception
    AILiarDataset(),
    PromptedSandbaggingDataset(),
    WerewolfPaired(),
]

prepared_datasets = prepare_datasets(datasets, tokenizer, model)  # type: ignore
# %%
method = "mms"
results = evaluate_datasets(prepared_datasets, prepared_insider_trading, method=method)
# %%
filename = f"results/generalization_results_{method}.csv"
save_generalization_results(results, filename)

filename = "results/generalization_results_mms.csv"
results = load_generalization_results(filename)
fig = create_plot(
    results, attribute="auroc", ordering=[d.name for d in datasets] + ["Insider Trading"]
)
fig.show()
# %%
