# %%
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import NamedTuple, cast

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from matplotlib.figure import Figure
from matplotlib.ticker import PercentFormatter

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import PromptScorer, create_prompt_scorer
from deception_detection.results.plot import get_fprs_tprs
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,
    to_readable_dataset_name,
)
from deception_detection.types import Label
from deception_detection.utils import REPO_ROOT

# %%

OUT_DIR = REPO_ROOT.parent / "dd_overleaf/figures"


@dataclass
class ProbeConfig:
    """Configuration for a single probe experiment."""

    name: str
    title: str
    exp_path: Path


class ExperimentInfo(NamedTuple):
    exp: Experiment
    scorer: PromptScorer
    eval_df: pd.DataFrame
    control_df: pd.DataFrame
    title: str

    @classmethod
    def from_exp(
        cls,
        exp: Experiment,
        title: str,
        exclude_datasets: set[str],
    ) -> "ExperimentInfo":
        scorer = create_prompt_scorer("mean", exp, fprs=[0.01])
        dataset_names = [name for name in exp.get_eval_scores() if name not in exclude_datasets]
        eval_df = get_eval_df(exp, scorer, dataset_names)
        control_df = get_control_df(exp, scorer)
        return cls(exp, scorer, eval_df, control_df, title)


def get_eval_df(exp: Experiment, scorer: PromptScorer, dataset_names: list[str]) -> pd.DataFrame:
    dfs: list[pd.DataFrame] = []
    for dataset_name in dataset_names:
        df = get_dataframe_of_scores_on_dataset(exp, dataset_name, scorer=scorer)
        df["dataset"] = to_readable_dataset_name(dataset_name, separate_val=True)
        dfs.append(df)
    return pd.concat(dfs, ignore_index=True)


def plot_grid_of_detailed_results(
    probe_configs: Sequence[ProbeConfig],
    exclude_datasets: set[str],
    output_path: Path | None = None,
    figsize: tuple[float, float] = (11, 12),
) -> Figure:
    """Create a grid of detailed results for multiple probe experiments.

    Args:
        probe_configs: Sequence of ProbeConfig objects defining the experiments to plot
        exclude_datasets: Set of dataset names to exclude from the analysis
        output_path: Optional path to save the figure
        figsize: Size of the figure in inches

    Returns:
        matplotlib.figure.Figure: The generated figure
    """
    # Create experiments from configs
    exps = [Experiment(ExperimentConfig.from_path(config.exp_path)) for config in probe_configs]

    exp_infos = [
        ExperimentInfo.from_exp(exp, config.title, exclude_datasets)
        for exp, config in zip(exps, probe_configs, strict=True)
    ]

    fig = plt.figure(figsize=figsize)
    gs = fig.add_gridspec(len(exp_infos), 2, hspace=0.4, width_ratios=[1, 0.8])

    for row, exp_info in enumerate(exp_infos):
        exp, scorer, eval_df, control_df, title = exp_info

        # Get the position of the current row's subplots
        row_bbox = gs[row, :].get_position(fig)
        # Add title just above the plots with a small offset
        fig.text(0.5, row_bbox.y1 + 0.01, title, ha="center", va="bottom", fontsize=9)

        # Violin plot takes up first column
        ax_violin = fig.add_subplot(gs[row, 0])

        # Create a nested gridspec for the ROC plots in the second column
        gs_roc = gs[row, 1].subgridspec(1, 2)
        ax_honest = fig.add_subplot(gs_roc[0, 0])
        ax_control = fig.add_subplot(gs_roc[0, 1], sharex=ax_honest)

        # Get data for each dataset
        dataset_names = [name for name in exp.get_eval_scores() if name not in exclude_datasets]

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

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

        # Plot datasets one by one
        for dataset in eval_df["dataset"].unique():
            dataset = cast(str, dataset)
            subset = eval_df[eval_df["dataset"] == dataset]
            sns.violinplot(
                data=subset,  # type: ignore
                x="dataset",
                order=[dataset],
                hue_order=["honest", "ambiguous", "deceptive"],
                ax=ax_violin,
                legend=dataset == dataset_names[0],  # only show legend for first dataset
                **args,  # type: ignore
            )

        # Plot control data
        sns.violinplot(data=control_df, x="dataset", ax=ax_violin, **args, legend=False)  # type: ignore

        ax_violin.set_xlabel("")
        ax_violin.set_ylabel("Probe deceptiveness score")
        ax_violin.set_yticks([])
        ax_violin.axhline(y=0, color="black", linewidth=1)

        for dataset_name, scores in exp.get_eval_scores().items():
            if dataset_name in exclude_datasets:
                continue

            split_scores = scores.get_scores_by_label()
            if len(split_scores[Label.HONEST]) == 0:
                continue

            honest = scorer.process(split_scores[Label.HONEST])
            deceptive = scorer.process(split_scores[Label.DECEPTIVE])
            control = scorer.process(exp.get_control_scores()["alpaca__plain"])

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

            # vs honest
            fprs, tprs = get_fprs_tprs(honest, deceptive)
            ax_honest.plot(fprs, tprs, color=color, label=name)

            # vs control
            fprs, tprs = get_fprs_tprs(control, deceptive)
            ax_control.plot(fprs, tprs, color=color, label=name)

        for ax, is_honest in [(ax_honest, True), (ax_control, False)]:
            ax.plot([0, 1], [0, 1], "k--", linewidth=1)
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
            ax.set_xlabel("Honest FPR" if is_honest else "Control FPR")
            if is_honest:
                ax.set_ylabel("Deceptive TPR")
                ax.yaxis.set_major_formatter(PercentFormatter(xmax=1, decimals=0))
            else:
                ax.set_yticks([])
            ax.xaxis.set_major_formatter(PercentFormatter(xmax=1, decimals=0))
            if row == 0 and not is_honest:
                ax.legend(loc="lower right")

            ax.set_box_aspect(1)

    plt.tight_layout()

    if output_path is not None:
        save_plot(output_path)
        plt.close(fig)

    return fig
