from pathlib import Path

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 PromptScorerMethod, create_prompt_scorer
from deception_detection.results.utils import save_plot, 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,
    separate_val_dataset_name: 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, separate_val_dataset_name)
                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}")
            raise 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 = torch.tensor(
                    exp.get_black_box_control_scores()["alpaca__plain"].scores
                )
                alpaca_threshold = torch.nanquantile(alpaca_scores, 0.99).item()
                print("Black box threshold", alpaca_threshold)
                frac_over_threshold = (alpaca_scores >= alpaca_threshold).float().mean().item()
                print(f"frac alpaca >= threshold {frac_over_threshold:.2%}")
                frac_strictly_over_threshold = (
                    (alpaca_scores > alpaca_threshold).float().mean().item()
                )
                print(f"frac alpaca > threshold {frac_strictly_over_threshold:.2%}")
                print("using strict threshold")
                is_over_threshold = (
                    torch.tensor(split_scores[Label.DECEPTIVE].scores) > alpaca_threshold
                )
                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,
    outfile: Path | None = None,
    separate_val_dataset_name: 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,
        separate_val_dataset_name=separate_val_dataset_name,
    )
    # Create the plot
    _, axs = plt.subplots(2, 1, figsize=figsize, sharex=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,
    }
    barplot = 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")
    if use_recall_for_control:
        axs[1].set_ylabel(r"Recall at 1% FPR")

    plt.ylim(*y_lim)
    axs[0].get_legend().remove()  # we need to make the legend to get handles but then remove it
    plt.xlabel("")
    handles, labels = barplot.get_legend_handles_labels()
    plt.figlegend(handles, labels, bbox_to_anchor=(1.0, 0.5), loc="center left")
    
    # Fix x-axis labels to prevent overlap
    for ax in axs:
        ax.tick_params(axis='x', rotation=45, labelsize=9)
        for label in ax.get_xticklabels():
            label.set_horizontalalignment('right')
    
    plt.tight_layout()
    # Add extra bottom margin for rotated labels
    plt.subplots_adjust(bottom=0.15)

    output_path = outfile or results_folder / "auroc_comparison.png"
    save_plot(output_path)


def make_reduction_comparison_bar_chart(
    results_folder: Path,
    score_methods: list[PromptScorerMethod],
    figsize: tuple[float, float] = (8, 5),
    excluded_datasets: set[str] | None = None,
    y_mins: tuple[float, float] = (0.5, 1.0),
    outfile: Path | None = None,
    relu_threshold: float | None = None,
    legend: bool = True,
) -> 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)
    """

    exp = load_experiment(results_folder)
    scorers = {
        scoring_method: create_prompt_scorer(scoring_method, exp, fprs=[0.01], default_fpr=0.01)
        for scoring_method in score_methods
    }
    if relu_threshold is not None:
        scorers["relu"].relu_threshold = relu_threshold  # type: ignore
    data = []
    all_dataset_scores: list[Scores] = []
    for scoring_method in score_methods:
        scorer = scorers[scoring_method]
        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)
            auroc_alpaca = scorer.recall_at_fpr(scores)
            readable_name = to_readable_dataset_name(dataset)
            data.append(
                {
                    "scoring_method": scoring_method,
                    "Dataset": readable_name,
                    "AUROC": auroc,
                    "AUROC vs Alpaca": auroc_alpaca,
                }
            )

    df = pd.DataFrame(data)
    # Create the plot
    plt.clf()
    _, axs = plt.subplots(2, 1, figsize=figsize, sharex=True)

    palette = {
        name: color for name, color in zip(score_methods, sns.color_palette("deep"), strict=False)
    }

    kwargs = {
        "data": df,
        "x": "Dataset",
        "hue": "scoring_method",
        "hue_order": score_methods,
        "palette": palette,
    }
    sns.barplot(**kwargs, y="AUROC", ax=axs[0], legend=legend)  # type: ignore
    sns.barplot(**kwargs, y="AUROC vs Alpaca", ax=axs[1], legend=False)  # type: ignore
    if legend:
        axs[0].legend(bbox_to_anchor=(1.02, 1), loc="upper left", title="Reduction")
    axs[1].set_ylabel(r"Recall at 1% FPR")

    axs[0].set_ylim(y_mins[0], 1.0)
    axs[1].set_ylim(y_mins[1], 1.0)
    axs[0].set_xlabel("")
    axs[1].set_xlabel("")

    # Fix x-axis labels to prevent overlap
    for ax in axs:
        ax.tick_params(axis='x', rotation=45, labelsize=9)
        for label in ax.get_xticklabels():
            label.set_horizontalalignment('right')

    plt.tight_layout()
    # Add extra bottom margin for rotated labels
    plt.subplots_adjust(bottom=0.15)
    output_path = outfile or results_folder / "reduction_comparison.png"
    save_plot(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)
