from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.results.utils import (
    get_dataframe_of_scores_on_dataset,
    to_readable_dataset_name,
)
from deception_detection.utils import REPO_ROOT


def load_experiment(
    folder: str | Path,
) -> Experiment:
    """Load an experiment from a folder."""
    return Experiment(ExperimentConfig.from_path(Path(folder)))


def get_combined_df(exp1: Experiment, exp2: Experiment, dataset_name: str) -> pd.DataFrame:
    scorer = create_prompt_scorer("mean", exp1, fprs=[0.01], default_fpr=0.01)
    df1 = get_dataframe_of_scores_on_dataset(exp1, dataset_name, scorer, normalize_scores=True)
    df2 = get_dataframe_of_scores_on_dataset(exp2, dataset_name, scorer, normalize_scores=True)
    return pd.merge(df1, df2, suffixes=("_1", "_2"), left_index=True, right_index=True)


vary_train_folder = REPO_ROOT / "results" / "comparisons"
exp1 = load_experiment(vary_train_folder / "mainline")
exp2 = load_experiment(vary_train_folder / "roleplaying")

EVAL_DATASETS = [
    "insider_trading__onpolicy",
    "insider_trading_doubledown__onpolicy",
    "sandbagging_v2__wmdp_mmlu",
    "ai_liar__original_without_answers",
]

# Create a 2x2 subplot
fig, axes = plt.subplots(2, 2, figsize=(8, 8), sharex=True, sharey=True)
axes = axes.flatten()

for idx, dataset in enumerate(EVAL_DATASETS):
    df = get_combined_df(exp1, exp2, dataset)
    df = df[df["label_1"].isin(["honest", "deceptive"])]

    ax = axes[idx]
    ax.scatter(
        df["score_1"],  # type: ignore
        df["score_2"],  # type: ignore
        c=["r" if label == "deceptive" else "b" for label in df["label_1"]],  # type: ignore
        alpha=0.5,
        s=2,
    )
    ax.set_xlabel("Mainline")
    ax.set_ylabel("Roleplaying")
    ax.set_title(to_readable_dataset_name(dataset))

# Add a legend
fig.legend(["Honest", "Deceptive"], loc="center right")
plt.tight_layout()  # Adjust spacing between subplots
plt.show()
