from pathlib import Path

import fire
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import torch

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.results.utils import to_readable_dataset_name
from deception_detection.scores import Scores
from deception_detection.types import Label


def load_experiment(experiment_path: Path) -> Experiment:
    return Experiment(ExperimentConfig.from_path(experiment_path))


def get_aurocs(exp: Experiment) -> dict[str, float]:
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01], default_fpr=0.01)
    return {
        dataset_name: scorer.auroc(scores) for dataset_name, scores in exp.get_eval_scores().items()
    }


def get_aruoc_vs_alpaca(exp: Experiment) -> dict[str, float]:
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01], default_fpr=0.01)
    return {
        dataset_name: scorer.auroc_alpaca(scores)
        for dataset_name, scores in exp.get_eval_scores().items()
    }


def get_aurocs_df(
    results_folder: Path,
    label_key: str,
    value_names: dict[str, str],
    excluded_datasets: set[str] | None = None,
    exp_for_black_box: Path | None = None,
    use_recall_for_control: bool = False,
    include_merged: bool = False,
) -> pd.DataFrame:
    data = []
    for exp_folder in results_folder.iterdir():
        if not exp_folder.is_dir():
            continue

        try:
            exp = load_experiment(exp_folder)
            exp_value_name = value_names[exp.cfg.model_dump()[label_key]]
            scorer = create_prompt_scorer("mean", exp, fprs=[0.01], default_fpr=0.01)
            all_dataset_scores: list[Scores] = []
            for dataset, scores in exp.get_eval_scores().items():
                if excluded_datasets and dataset in excluded_datasets:
                    continue
                all_dataset_scores.append(scores)

                auroc = scorer.auroc(scores)
                if use_recall_for_control:
                    auroc_alpaca = scorer.recall_at_fpr(scores)
                else:
                    auroc_alpaca = scorer.auroc_alpaca(scores)
                readable_name = to_readable_dataset_name(dataset)
                data.append(
                    {
                        label_key: exp_value_name,
                        "Dataset": readable_name,
                        "AUROC": auroc,
                        "AUROC vs Alpaca": auroc_alpaca,
                    }
                )

            if include_merged:
                merged_scores = Scores.from_list(all_dataset_scores)
                lens = [
                    len(scores.get_scores_by_label()[Label.HONEST])
                    + len(scores.get_scores_by_label()[Label.DECEPTIVE])
                    for scores in all_dataset_scores
                ]
                weights = torch.cat([torch.ones(dataset_len) / dataset_len for dataset_len in lens])
                auroc_across_datasets = scorer.auroc(merged_scores, weights=weights)
                if use_recall_for_control:
                    auroc_across_datasets_alpaca = scorer.recall_at_fpr(merged_scores)
                else:
                    auroc_across_datasets_alpaca = scorer.auroc_alpaca(merged_scores)
                data.append(
                    {
                        label_key: exp_value_name,
                        "Dataset": "Merged",
                        "AUROC": auroc_across_datasets,
                        "AUROC vs Alpaca": auroc_across_datasets_alpaca,
                    }
                )

        except Exception as e:
            print(f"Failed to load results from {exp_folder}: {e}")
            continue

    if exp_for_black_box:
        exp = load_experiment(exp_for_black_box)
        black_box_scores = exp.get_black_box_scores()
        black_box_scores_list: list[Scores] = []
        scorer = create_prompt_scorer("mean", exp, fprs=[0.01], default_fpr=0.01)
        for dataset, scores in black_box_scores.items():
            if excluded_datasets and dataset in excluded_datasets:
                continue
            black_box_scores_list.append(scores)
            auroc = scorer.auroc(scores)
            if auroc == float("nan"):
                print(f"No black box scores found for dataset {dataset}")

            if use_recall_for_control:
                split_scores = scores.get_scores_by_label()
                alpaca_scores = exp.get_black_box_control_scores()["alpaca__plain"]
                alpaca_threshold = torch.quantile(torch.tensor(alpaca_scores.scores), 0.99).item()
                # generally alpaca_threshold is 6
                # It's not clear which of >= or > is more principled (do you want FPR~0.3% or ~6%?)
                # but the choice does significantly change the results.
                is_over_threshold = (
                    torch.tensor(split_scores[Label.DECEPTIVE].scores) >= alpaca_threshold
                )
                print(alpaca_threshold, is_over_threshold.float().mean().item())
                auroc_alpaca = is_over_threshold.float().mean().item()
            else:
                auroc_alpaca = float("nan")
            data.append(
                {
                    label_key: value_names["black_box"],
                    "Dataset": to_readable_dataset_name(dataset),
                    "AUROC": auroc,
                    "AUROC vs Alpaca": auroc_alpaca,
                }
            )

        if include_merged:
            merged_scores = Scores.from_list(black_box_scores_list)
            auroc_across_datasets = scorer.auroc(merged_scores)
            data.append(
                {
                    label_key: value_names["black_box"],
                    "Dataset": "Merged",
                    "AUROC": auroc_across_datasets,
                    "AUROC vs Alpaca": float("nan"),
                }
            )

    df = pd.DataFrame(data)
    print(df)
    return df


def make_comparison_bar_chart(
    results_folder: Path,
    label_key: str,
    value_names: dict[str, str],
    exp_for_black_box: Path | None = None,
    figsize: tuple[float, float] = (8, 5),
    excluded_datasets: set[str] | None = None,
    y_lim: tuple[float, float] = (0.5, 1.0),
    use_recall_for_control: bool = False,
    legend_title: str | None = None,
    include_merged: bool = False,
) -> None:
    """
        Create a bar chart comparing AUROCs across different experiments.
    a
        Args:
            results_folder: Path to folder containing experiment results
            label_key: Key in the experiment config to use for labeling experiments
            value_names: Dictionary mapping experiment config values to names for the legend
            figsize: Size of the figure (width, height)
    """

    df = get_aurocs_df(
        results_folder=results_folder,
        label_key=label_key,
        value_names=value_names,
        excluded_datasets=excluded_datasets,
        exp_for_black_box=exp_for_black_box,
        use_recall_for_control=use_recall_for_control,
        include_merged=include_merged,
    )
    # Create the plot
    _, axs = plt.subplots(2, 1, figsize=figsize, sharey=True)

    palette = {
        name: color
        for name, color in zip(value_names.values(), sns.color_palette("deep"), strict=False)
    }
    if "black_box" in value_names:
        palette[value_names["black_box"]] = (0.4, 0.4, 0.4)

    kwargs = {
        "data": df,
        "x": "Dataset",
        "hue": label_key,
        "hue_order": list(value_names.values()),
        "palette": palette,
    }
    sns.barplot(**kwargs, y="AUROC", ax=axs[0], legend=True)  # type: ignore
    sns.barplot(**kwargs, y="AUROC vs Alpaca", ax=axs[1], legend=False)  # type: ignore
    axs[0].set_title("AUROC vs Honest")
    if use_recall_for_control:
        axs[1].set_title("Recall vs Control")
    else:
        axs[1].set_title("AUROC vs Control")
    axs[0].legend(bbox_to_anchor=(1.02, 1), loc="upper left", title=legend_title)
    if use_recall_for_control:
        axs[1].set_ylabel(r"Recall at 1% FPR")

    plt.ylim(*y_lim)
    axs[0].set_xlabel("")
    axs[1].set_xlabel("")

    plt.tight_layout()
    output_path = results_folder / "auroc_comparison.png"
    plt.savefig(output_path)


def make_comparison_point_plot(
    results_folder: Path,
    label_key: str,
    value_names: dict[str, str],
    excluded_datasets: set[str] | None = None,
) -> None:
    results_folder = Path(results_folder)
    df = get_aurocs_df(results_folder, label_key, value_names, excluded_datasets)

    plt.figure(figsize=(9, 3))
    kwargs = {
        "data": df,
        "x": "Dataset",
        "hue": label_key,
        "hue_order": list(value_names.values()),
        "palette": "deep",
        "dodge": True,
    }

    g1 = sns.stripplot(**kwargs, y="AUROC", marker="o", size=6)  # type: ignore
    sns.stripplot(**kwargs, y="AUROC vs Alpaca", marker="*", size=10)  # type: ignore

    # Create custom legend handles
    color_handles = [
        plt.Line2D([], [], marker="s", color=h.get_color(), linestyle="None")  # type: ignore
        for h in g1.legend_.legend_handles[: len(value_names)]  # type: ignore
    ]
    shape_handles = [
        plt.Line2D([], [], marker="o", color="gray", linestyle="None", label="vs Honest"),
        plt.Line2D(
            [], [], marker="*", markersize=10, color="gray", linestyle="None", label="vs Control"
        ),
    ]
    legend_labels = ["vs Honest", "vs Control"] + list(value_names.values())

    # Combine all handles and labels
    plt.legend(
        handles=shape_handles + color_handles,
        labels=legend_labels,
        bbox_to_anchor=(1.03, 0.5),
        loc="center left",
    )

    plt.ylabel("AUROC")
    plt.ylim(0.65, 1.0)
    plt.xlabel("")

    plt.tight_layout()
    output_path = results_folder / "auroc_comparison_point.png"
    plt.savefig(output_path)


def rp_off_policy_plot():
    make_comparison_bar_chart(
        results_folder=Path("results/rp_offpolicy"),
        label_key="name",
        value_names={
            "mainline": "Mainline (RepE)",
            "on_policy": "On Policy",
            "off_policy": "Off Policy",
            "mean_lr": "Mean before LR",
            "last_token": "Last Token",
        },
        excluded_datasets={
            "roleplaying__offpolicy_train",
            "repe_honesty__plain",
            "ai_audit__no_reasoning",
        },
    )


def comparison_plot():
    make_comparison_bar_chart(
        results_folder=Path("results/comparisons"),
        label_key="name",
        # legend_title="Probe training set",
        value_names={
            "mainline": "Train on Instructed-Pairs",
            "sae_repe": "Train on Instructed-Pairs\nSAE activations",
            "roleplaying": "Train on Roleplaying",
            "sae": "Train on Roleplaying\nSAE activations",
            "descriptive": "Train on single\nDescriptive Question",
            "black_box": "Black-Box baseline",
        },
        excluded_datasets={
            "ai_audit__no_reasoning",
            "repe_honesty__plain",
            "ai_liar__original_without_answers",
        },
        exp_for_black_box=Path("results/comparisons/mainline"),
        y_lim=(0.0, 1.0),
        figsize=(10, 5),
        use_recall_for_control=True,
        # include_merged=True,
    )


def method_comparison_plot():
    make_comparison_bar_chart(
        results_folder=Path("results/by_method"),
        label_key="name",
        value_names={
            "lr": "Logistic Regression",
            "mms": "Difference of Means",
            "cmms": "Linear Discriminant Analysis",
            "lat": "Linear Artificial Tomography",
        },
        excluded_datasets={
            "ai_audit__no_reasoning",
            "repe_honesty__plain",
            "ai_liar__original_without_answers",
        },
        y_lim=(0.0, 1.0),
        figsize=(10, 5),
        use_recall_for_control=True,
    )


def only_eot_comparison():
    make_comparison_bar_chart(
        results_folder=Path("results/vary_only_eot"),
        label_key="name",
        value_names={
            "roleplaying": "entire response",
            "roleplaying_only_eot": "last token",
        },
        use_recall_for_control=True,
        excluded_datasets={"ai_liar__original_without_answers"},
    )


if __name__ == "__main__":
    fire.Fire(
        {
            "plot": make_comparison_bar_chart,
            "rp_off_policy": rp_off_policy_plot,
            "comparison": comparison_plot,
            "method": method_comparison_plot,
            "eot": only_eot_comparison,
        }
    )
