from pathlib import Path
from typing import Literal, cast

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import torch
from matplotlib.patches import Rectangle
from matplotlib.ticker import PercentFormatter
from sklearn.metrics import auc
from torch import Tensor

from deception_detection.experiment import Experiment
from deception_detection.log import logger
from deception_detection.metrics import get_auroc
from deception_detection.prompt_scorer import PromptScorer
from deception_detection.results.utils import (
    AMBIGUOUS_COLOR,
    DECEPTIVE_COLOR,
    HONEST_COLOR,
    NEUTRAL_COLOR,
    get_control_df,
    get_dataframe_of_scores_on_dataset,
    get_dataset_color,
    save_plot_exp,
    to_readable_dataset_name,
)
from deception_detection.scores import Scores
from deception_detection.types import Label


def get_fprs_tprs(honest: Tensor, deceptive: Tensor):
    min_score = min(honest.min().item(), deceptive.min().item())
    max_score = max(honest.max().item(), deceptive.max().item())
    thresholds = torch.linspace(min_score - 1, max_score, 1000)

    def positive_rates(scores: Tensor, thresholds: Tensor) -> list[float]:
        return (scores > thresholds[:, None]).float().mean(dim=-1).tolist()

    tprs = positive_rates(deceptive, thresholds)
    fprs = positive_rates(honest, thresholds)
    return fprs, tprs


def violin_plot(
    experiment: Experiment,
    scorer: PromptScorer,
    include_train: bool = True,
    exclude_datasets: set[str] | None = None,
    legend_loc: str = "lower right",
    outfile: Path | None = None,
    figsize: tuple[float, float] | None = None,
):
    normalize = True
    dfs: list[pd.DataFrame] = []
    if exclude_datasets is None:
        exclude_datasets = {"ai_liar__original_without_answers", "repe_honesty__plain_with_sys_val"}
    dataset_names = [
        name
        for name in experiment.get_eval_scores()
        if name not in exclude_datasets
        and (include_train or not name.startswith(experiment.cfg.train_data))
    ]
    for dataset_name in dataset_names:
        df = get_dataframe_of_scores_on_dataset(
            experiment, dataset_name, scorer=scorer, normalize_scores=normalize
        )
        df["dataset"] = to_readable_dataset_name(dataset_name, separate_val=True)
        dfs.append(df)
    eval_df = pd.concat(dfs, ignore_index=True)
    control_df = get_control_df(experiment, scorer, normalize)

    # Create seaborn violin plot
    figsize = figsize or (1 + 1.4 * len(dataset_names), 3.5)
    _, axs = plt.subplots(
        1, 2, figsize=figsize, sharey=True, width_ratios=[len(dataset_names), 0.6]
    )

    palette = {
        "honest": HONEST_COLOR,
        "deceptive": DECEPTIVE_COLOR,
        "control": NEUTRAL_COLOR,
        "ambiguous": AMBIGUOUS_COLOR,
    }

    def get_order(subset: pd.DataFrame) -> list[str]:
        if "ambiguous" in set(subset["label"].unique()):  # type: ignore
            return ["honest", "ambiguous", "deceptive"]
        else:
            return ["honest", "deceptive"]

    args = dict(y="score", hue="label", cut=0, inner=None, palette=palette)

    # plot datasets 1 by 1 to allow for different `hue_order`
    for dataset in cast(list[str], eval_df["dataset"].unique().tolist()):
        subset = cast(pd.DataFrame, eval_df[eval_df["dataset"] == dataset])  # type: ignore[reportUnnecessaryCast]
        sns.violinplot(
            data=subset,
            x="dataset",
            order=[dataset],
            hue_order=get_order(subset),
            ax=axs[0],
            legend=dataset == "Insider Trading\nReport",  # only one legend
            **args,  # type: ignore
        )

    sns.violinplot(data=control_df, x="dataset", ax=axs[1], **args, legend=False)  # type: ignore

    # plt.suptitle("Deceptiveness scores by honesty")
    axs[0].set_ylabel("Deceptiveness score")
    axs[0].set_yticks([])
    axs[1].set_ylabel("")
    axs[1].set_xlabel("")
    axs[0].set_xlabel("")
    axs[0].legend(title="", loc=legend_loc)
    threshold = 0 if normalize else scorer.prompt_score_threshold(scorer.default_fpr)
    axs[0].axhline(y=threshold, color="black", linewidth=1)
    axs[1].axhline(y=threshold, color="black", linewidth=1)

    # Fix x-axis label formatting to prevent overlap
    for ax in axs:
        ax.tick_params(axis='x', rotation=45, labelsize=8)
        # Add some padding to bottom margin for rotated labels
        for label in ax.get_xticklabels():
            label.set_horizontalalignment('right')

    # Save plot with adjusted layout
    plt.tight_layout()
    # Add extra bottom margin for rotated labels
    plt.subplots_adjust(bottom=0.2)
    save_plot_exp(experiment, f"violin_plot_{scorer.method}.png", outfile)


def plot_roc_comparison(
    experiment: Experiment,
    scorer: PromptScorer,
    exclude_datasets: set[str] | None = None,
    outfile: Path | None = None,
    figsize: tuple[float, float] | None = None,
    x_maxs: tuple[float, float] = (1, 0.03),
    x_line: float | None = 0.01,
    control_dataset_name: str = "alpaca__plain",
):
    fig, axs = plt.subplots(1, 2, figsize=figsize or (8, 4), sharey=True)

    all_scores = experiment.get_eval_scores()
    alpaca_scores = scorer.process(experiment.get_control_scores()[f"{control_dataset_name}"])

    for dataset_name, scores in all_scores.items():
        if exclude_datasets is not None and dataset_name in exclude_datasets:
            continue
        split_scores = scores.get_scores_by_label()
        if len(split_scores[Label.HONEST]) == 0:
            logger.warning(f"no honest scores for {dataset_name}")
            continue
        honest = scorer.process(split_scores[Label.HONEST])
        deceptive = scorer.process(split_scores[Label.DECEPTIVE])
        name = to_readable_dataset_name(dataset_name, separate_val=True)
        color = get_dataset_color(dataset_name)
        axs[0].plot(*get_fprs_tprs(honest, deceptive), label=name, color=color)
        axs[1].plot(*get_fprs_tprs(alpaca_scores, deceptive), label=name, color=color)

    axs[0].set_xlabel("FPR on Honest")
    axs[1].set_xlabel("FPR on Control")

    axs[0].set_ylabel("Recall (TPR)")
    axs[0].set_xlim(0, x_maxs[0])
    axs[1].set_xlim(0, x_maxs[1])
    axs[0].set_ylim(0, 1)
    axs[0].set_box_aspect(1)
    axs[1].set_box_aspect(1)

    # Format ticks as percentages
    axs[0].xaxis.set_major_formatter(
        PercentFormatter(xmax=1, decimals=1 if x_maxs[0] < 0.05 else 0)
    )
    axs[0].yaxis.set_major_formatter(PercentFormatter(xmax=1, decimals=0))
    axs[1].xaxis.set_major_formatter(
        PercentFormatter(xmax=1, decimals=1 if x_maxs[1] < 0.05 else 0)
    )
    axs[1].yaxis.set_major_formatter(PercentFormatter(xmax=1, decimals=0))

    axs[0].plot([0, 1], [0, 1], color="k", linestyle="--", linewidth=1)
    axs[1].plot([0, 1], [0, 1], color="k", linestyle="--", linewidth=1)
    if x_line:
        axs[1].axvline(x_line, color="k", linestyle="-", linewidth=1)

    axs[1].legend(loc="lower right")

    save_plot_exp(experiment, f"all_roc_plots_{scorer.method}.png", outfile)
    plt.close(fig)


def roleplaying_violin(
    exp: Experiment,
    scorer: PromptScorer,
    outfile: Path | None = None,
    figsize: tuple[float, float] = (5, 4),
):
    """Violin plot of roleplaying scores by deceptive grade, with alpaca scores as control"""
    df = get_dataframe_of_scores_on_dataset(exp, "roleplaying__plain", scorer=scorer)
    control_df = get_control_df(exp, scorer)
    control_df["model_graded_deception"] = "Control"
    combined_df = pd.concat([df, control_df])

    plt.figure(figsize=figsize)
    palette = [HONEST_COLOR] * 2 + [AMBIGUOUS_COLOR] * 3 + [DECEPTIVE_COLOR] * 2 + [NEUTRAL_COLOR]
    sns.violinplot(
        data=combined_df,
        x="model_graded_deception",
        y="score",
        order=[1, 2, 3, 4, 5, 6, 7, "Control"],
        hue_order=[1, 2, 3, 4, 5, 6, 7, "Control"],
        palette=palette,
        hue="model_graded_deception",
        legend=False,
        cut=0,
    )
    # plt.title("Roleplaying vs Alpaca")
    plt.xlabel("Model Graded Deception")
    plt.ylabel("Probe Deceptiveness Score")
    plt.yticks([])

    # compute auroc reclassifying ambiguous
    honest_scores = df[df["model_graded_deception"] <= 4]["score"]
    deceptive_scores = df[df["model_graded_deception"] > 4]["score"]
    auroc = get_auroc(honest_scores, deceptive_scores)  # type: ignore
    print(f"AUROC reclassifying ambiguous: {auroc:.3f}")
    threshold = np.quantile(control_df["score"], 0.99)  # type: ignore
    recall_at_threshold = (deceptive_scores > threshold).mean()
    print(f"Recall at threshold: {recall_at_threshold:.3f}")

    save_plot_exp(exp, f"roleplaying_{scorer.method}.png", outfile)


def ai_audit_violin(
    exp: Experiment,
    scorer: PromptScorer,
    outfile: Path | None = None,
    figsize: tuple[float, float] = (5, 2.5),
) -> None:
    """Create violin plot for AI Audit data, separated by experiment and variation."""
    normalize = True
    df = get_dataframe_of_scores_on_dataset(exp, "ai_audit__a1_2_reasoning", scorer, normalize)
    df["label"] = df["label"].map(lambda label: label.title())  # type: ignore
    control_df = get_control_df(exp, scorer)
    control_df["label"] = "Control"
    df = pd.concat([df, control_df])

    plt.figure(figsize=figsize)
    sns.violinplot(
        data=df,
        y="score",
        hue="label",
        cut=0,
        palette=[HONEST_COLOR, DECEPTIVE_COLOR, NEUTRAL_COLOR],
        hue_order=["Honest", "Deceptive", "Control"],
        inner=None,
    )
    plt.yticks([])
    plt.xticks([])
    plt.ylabel("Probe deceptiveness score")
    plt.legend(title="", loc="upper left")

    scores = exp.get_eval_scores()["ai_audit__a1_2_reasoning"]
    auroc = scorer.auroc(scores)
    recall_at_threshold = scorer.recall_at_fpr(scores, 0.01)
    plt.axhline(y=0, color="k", linewidth=1)
    # annotate with auroc in lower left
    plt.text(0.05, 0.05, f"AUROC: {auroc:.3f}", transform=plt.gca().transAxes)
    plt.text(0.05, 0.15, f"Recall @1%: {recall_at_threshold:.3f}", transform=plt.gca().transAxes)
    plt.tight_layout()
    save_plot_exp(exp, f"ai_audit_violin_{scorer.method}.png", outfile)


def plot_roc_curve_with_auroc(
    experiment: Experiment,
    scorer: PromptScorer,
    compare_to: Literal["honest", "control"] = "honest",
    exclude_datasets: set[str] | None = None,
    x_max: float = 1.0,
    include_legend: bool = False,
    include_title: bool = False,
    legend_outside: bool = True,
    legend_with_auc: bool = True,
    figsize: tuple[float, float] = (3.5, 3.5),
    xline: float | None = 0.01,
    outfile: Path | None = None,
    control_dataset_name: str = "alpaca__plain",
) -> None:
    """Plot ROC curves for deceptive vs honest/control with AUROC values in legend.

    Args:
        experiment: The experiment containing the scores
        scorer: The scorer to process the scores
        compare_to: Whether to compare deceptive examples against honest or control examples
        exclude_datasets: Optional set of dataset names to exclude
    """
    plt.figure(figsize=figsize)
    ax = plt.gca()

    if legend_outside:
        # Move the main axis to make room for legend
        box = ax.get_position()
        ax.set_position((box.x0, box.y0, box.width * 0.8, box.height))

    exclude_datasets = exclude_datasets or set()
    legend_elements = []

    # Get control scores if needed
    control_scores = None
    if compare_to == "control":
        control_scores = scorer.process(experiment.get_control_scores()[f"{control_dataset_name}"])

    for dataset_name, scores in experiment.get_eval_scores().items():
        if dataset_name in exclude_datasets or dataset_name.startswith(experiment.cfg.train_data):
            continue

        split_scores = scores.get_scores_by_label()
        deceptive = scorer.process(split_scores[Label.DECEPTIVE])

        if compare_to == "honest":
            if len(split_scores[Label.HONEST]) == 0:
                logger.warning(f"no honest scores for {dataset_name}")
                continue
            comparison_scores = scorer.process(split_scores[Label.HONEST])
        else:
            assert control_scores is not None
            comparison_scores = control_scores

        fprs, tprs = get_fprs_tprs(comparison_scores, deceptive)
        auroc = auc(fprs, tprs)

        name = to_readable_dataset_name(dataset_name)
        color = get_dataset_color(dataset_name)

        plt.plot(fprs, tprs, color=color)[0]
        rect = Rectangle((0, 0), 1, 1, fc=color)
        auroc_str = f"{auroc:.3f}" if auroc > 0.99 else f"{auroc:.2f}"
        if legend_with_auc:
            legend_elements.append((rect, f"{name} ({auroc_str})"))
        else:
            legend_elements.append((rect, name))

    plt.plot([0, x_max], [0, x_max], color="k", linestyle="--", linewidth=1)
    plt.xlim(0, x_max)
    plt.ylim(0, 1)
    plt.gca().set_box_aspect(1)
    plt.xlabel(f"FPR on {'Honest' if compare_to == 'honest' else 'Control'} (%)")
    plt.ylabel("Recall (TPR) (%)")
    if include_title:
        plt.title(f"ROC Curves: Deceptive vs {'Honest' if compare_to == 'honest' else 'Control'}")

    if xline:
        plt.axvline(xline, color="k", linestyle="-", linewidth=1)

    if include_legend:
        # Add legend to the right of the plot
        plt.legend(
            [elem[0] for elem in legend_elements],
            [elem[1] for elem in legend_elements],
            loc="upper left" if legend_outside else "lower right",
            bbox_to_anchor=(1.02, 1) if legend_outside else None,
            title="Eval. Dataset",
        )

    comparison_type = "honest" if compare_to == "honest" else "control"
    save_plot_exp(experiment, f"roc_with_auroc_{comparison_type}_{scorer.method}.png", outfile)
    plt.close()


def ai_liar_roc_plot(
    exp: Experiment,
    scorer: PromptScorer,
    outfile: Path | None = None,
    figsize: tuple[float, float] = (7, 3.5),
    label_text_fontsize: float = 10,
):
    fig, axs = plt.subplots(1, 2, figsize=figsize, sharey=True)

    dataset_name = "ai_liar__original_without_answers"
    alpaca_scores = scorer.process(exp.get_control_scores()["alpaca__plain"])
    scores = exp.get_eval_scores()[dataset_name]
    split_scores = scores.get_scores_by_label()
    if len(split_scores[Label.HONEST]) == 0:
        logger.warning(f"no honest scores for {dataset_name}")
        return
    honest = scorer.process(split_scores[Label.HONEST])
    deceptive = scorer.process(split_scores[Label.DECEPTIVE])
    name = to_readable_dataset_name(dataset_name)
    color = get_dataset_color(dataset_name)
    axs[0].plot(*get_fprs_tprs(honest, deceptive), label=name, color=color)
    axs[1].plot(*get_fprs_tprs(alpaca_scores, deceptive), label=name, color=color)

    # axs[0].set_title("ROC Curve vs Honest")
    axs[0].set_xlabel("FPR on Honest")
    axs[0].set_ylabel("Recall (TPR)")
    axs[0].set_xlim(0, 1)
    axs[0].set_ylim(0, 1)
    axs[0].set_box_aspect(1)
    axs[0].plot([0, 1], [0, 1], color="k", linestyle="--", linewidth=1)

    # axs[1].set_title("ROC Curve vs Alpaca")
    axs[1].set_xlabel("FPR on Alpaca")
    axs[1].set_box_aspect(1)

    axs[1].set_xlim(0, 0.03)
    axs[1].plot([0, 1], [0, 1], color="k", linestyle="--", linewidth=1)
    axs[1].axvline(0.01, color="k", linestyle="-", linewidth=1)
    axs[1].set_xticks([0, 0.01, 0.02, 0.03])

    # annotate with auc
    fprs, tprs = get_fprs_tprs(honest, deceptive)
    auroc = auc(fprs, tprs)
    text_args = dict(
        fontsize=label_text_fontsize,
        ha="right",
    )
    axs[0].text(0.95, 0.07, f"AUROC: {auroc:.3f}", transform=axs[0].transAxes, **text_args)

    fprs, tprs = get_fprs_tprs(alpaca_scores, deceptive)
    auroc = auc(fprs, tprs)
    import numpy as np

    fpr_at_01_idx = np.argmin(np.abs(np.array(fprs) - 0.01))
    tpr_at_01 = tprs[fpr_at_01_idx]
    axs[1].text(0.95, 0.15, f"Recall @1%: {tpr_at_01:.3f}", transform=axs[1].transAxes, **text_args)
    axs[1].text(0.95, 0.07, f"AUROC: {auroc:.3f}", transform=axs[1].transAxes, **text_args)

    plt.tight_layout()
    save_plot_exp(exp, "ai_liar_roc_plot.png", outfile)
    plt.close(fig)


def ai_liar_violin(
    exp: Experiment,
    scorer: PromptScorer,
    outfile: Path | None = None,
    figsize: tuple[float, float] = (6, 4),
    show_yticks: bool = False,
    title: str = "Roleplaying prompts from Pacchiardi et al. (2023)",
):
    df = get_dataframe_of_scores_on_dataset(exp, "ai_liar__original_without_answers", scorer)

    def get_category(row: pd.Series) -> str:  # type: ignore
        if row["label"] == "honest" and row["expected_label"] == "honest":
            return "honest, honest"
        elif row["label"] == "honest" and row["expected_label"] == "deceptive":
            return "honest, deceptive"
        elif row["label"] == "deceptive":
            return "deceptive"
        elif row["label"] == "skip" or row["label"] == "ambiguous":
            return "skip"
        raise AssertionError(row)

    df["category"] = df.apply(get_category, axis=1)  # type: ignore

    def get_label(category: str) -> str:
        counts = len(df[df["category"] == category])  # type: ignore
        return {
            "honest, honest": "honest response to\n honest prompt",
            "honest, deceptive": "honest response to\n deceptive prompt",
            "deceptive": "deceptive response",
            "skip": "skipped",
        }[category] + f"\n(n={counts})"

    plt.figure(figsize=figsize)
    sns.violinplot(
        data=df,
        x="category",
        hue="category",
        y="score",
        order=[
            "honest, honest",
            "honest, deceptive",
            "deceptive",
        ],
        palette=[HONEST_COLOR, HONEST_COLOR, DECEPTIVE_COLOR],
        legend=False,
        formatter=get_label,
    )
    if title:
        plt.title(title)
    plt.ylabel("Probe deceptiveness score")
    plt.xlabel("")
    if not show_yticks:
        plt.yticks([])
    plt.tight_layout()
    plt.axhline(0, color="k", linewidth=1)
    save_plot_exp(exp, "ai_liar_violin.png", outfile)
    plt.close()


def reclassify_ambigious(exp: Experiment, scorer: PromptScorer):
    recalssify_as = {
        "insider_trading__upscale": Label.HONEST,
        "insider_trading_doubledown__upscale": Label.HONEST,
        "sandbagging_v2__wmdp_mmlu": Label.DECEPTIVE,
    }

    eval_scores = exp.get_eval_scores()
    for dataset_name, label in recalssify_as.items():
        # Skip if dataset is not in the current experiment
        if dataset_name not in eval_scores:
            continue
            
        by_orig_label = eval_scores[dataset_name].get_scores_by_label()
        new_ambiguous = by_orig_label[Label.AMBIGUOUS]
        new_ambiguous.labels = [label] * len(new_ambiguous)
        new_scores = Scores.from_list(
            [by_orig_label[Label.HONEST], by_orig_label[Label.DECEPTIVE], new_ambiguous]
        )

        auroc = scorer.auroc(new_scores)
        print(f"{dataset_name} reclassify ambigious auroc: {auroc:.3f}")
        recall = scorer.recall_at_fpr(new_scores, 0.01)
        print(f"{dataset_name} reclassify ambigious recall: {recall:.3f}")


def make_dataset_specific_plots(
    exp: Experiment,
    scorer: PromptScorer,
    outfile: Path | None = None,
):
    if "ai_liar__original_without_answers" in exp.get_eval_scores():
        ai_liar_roc_plot(exp, scorer, outfile)
        ai_liar_violin(exp, scorer, outfile)
    if "ai_audit__a1_2reasoning" in exp.get_eval_scores():
        ai_audit_violin(exp, scorer, outfile)
    if "roleplaying__plain" in exp.get_eval_scores():
        roleplaying_violin(exp, scorer, outfile)
    reclassify_ambigious(exp, scorer)
