"""Generates all experiment-based plots for the paper."""

import fire
import matplotlib

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.results.comparison_bar_chart import (
    make_comparison_bar_chart,
    make_reduction_comparison_bar_chart,
)
from deception_detection.results.plot import (
    ai_audit_violin,
    ai_liar_roc_plot,
    ai_liar_violin,
    plot_roc_comparison,
    roleplaying_violin,
    violin_plot,
)
from deception_detection.results.plot_grid_of_detailed_results import (
    ProbeConfig,
    plot_grid_of_detailed_results,
)
from deception_detection.scripts.by_reg_coeff import plot as plot_reg_sweep
from deception_detection.utils import REPO_ROOT

OUT_DIR = REPO_ROOT / "results/_plots"
# OUT_DIR = REPO_ROO / "dd_overleaf/figures"
RESULTS_DIR = REPO_ROOT / "example_results"


# %%

matplotlib.rcParams.update({"font.size": 7.5})


def repe_plots() -> None:
    """The 'mainline' repe (aka InstructedPairs) experiment.

    Saves:
    - violin plot
    - roc deceptive vs honest plot with AUROC values in legend
    - roc deceptive vs control plot. No AUROC values in legend, they are in separate table.
    - roc curves on Pacchiardi RP
    - violin plot on Pacchiardi RP
    """
    exp = Experiment(ExperimentConfig.from_path(RESULTS_DIR / "instructed_pairs"))
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01])
    exclude_datasets = {
        "repe_honesty__plain_val",
        "ai_liar__original_without_answers",
        "ai_audit__no_reasoning",
        "ai_audit__reasoning",
        "ai_audit__a1_2_reasoning",
        "ai_audit__a1_2",
    }
    folder = OUT_DIR / "probes/repe"
    folder.mkdir(parents=True, exist_ok=True)
    violin_plot(
        exp,
        scorer,
        exclude_datasets=exclude_datasets,
        outfile=folder / "violin.png",
        figsize=(4.5, 2.5),
    )
    plot_roc_comparison(
        exp,
        scorer,
        exclude_datasets=exclude_datasets,
        outfile=folder / "roc_both.png",
        figsize=(4.5, 2.5),
        x_maxs=(0.2, 0.015),
    )
    ai_liar_violin(
        exp,
        scorer,
        outfile=folder / "ai_liar_violin.png",
        show_yticks=False,
        figsize=(4, 2.5),
        title="",
    )
    ai_liar_roc_plot(
        exp,
        scorer,
        outfile=folder / "ai_liar_roc.png",
        figsize=(5, 2.5),
        label_text_fontsize=7.5,
    )
    roleplaying_violin(exp, scorer, outfile=folder / "roleplaying_violin.png", figsize=(4.5, 3))
    ai_audit_violin(exp, scorer, outfile=folder / "ai_audit_violin.png", figsize=(3, 2.5))


def roleplaying_plots() -> None:
    exp = Experiment(ExperimentConfig.from_path(RESULTS_DIR / "roleplaying"))
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01])
    folder = OUT_DIR / "probes/roleplaying"
    folder.mkdir(parents=True, exist_ok=True)
    exclude_datasets = {"repe_honesty__plain_val", "ai_liar__original_without_answers"}
    violin_plot(
        exp,
        scorer,
        include_train=True,
        exclude_datasets=exclude_datasets,
        outfile=folder / "violin.png",
        figsize=(4.5, 2.5),
    )
    plot_roc_comparison(
        exp,
        scorer,
        exclude_datasets=exclude_datasets,
        outfile=folder / "roc_both.png",
        figsize=(4.5, 2.5),
        x_maxs=(1.0, 1.0),
        x_line=None,
    )
    ai_liar_violin(
        exp,
        scorer,
        outfile=folder / "ai_liar_violin.png",
        show_yticks=False,
        figsize=(4, 2.5),
        title="",
    )


def sae_rp_plots() -> None:
    exp = Experiment(ExperimentConfig.from_path(RESULTS_DIR / "sae_rp"))
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01])
    folder = OUT_DIR / "probes/sae_rp"
    folder.mkdir(parents=True, exist_ok=True)
    exclude_datasets = {"repe_honesty__plain_val", "ai_liar__original_without_answers"}
    violin_plot(
        exp,
        scorer,
        include_train=True,
        exclude_datasets=exclude_datasets,
        outfile=folder / "violin.png",
        figsize=(4.5, 2.5),
    )
    plot_roc_comparison(
        exp,
        scorer,
        exclude_datasets=exclude_datasets,
        outfile=folder / "roc_both.png",
        figsize=(4.5, 2.5),
        x_maxs=(1.0, 1.0),
        x_line=None,
    )


def descriptive_plots() -> None:
    exp = Experiment(ExperimentConfig.from_path(RESULTS_DIR / "descriptive"))
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01])
    folder = OUT_DIR / "probes/descriptive"
    folder.mkdir(parents=True, exist_ok=True)
    exclude_datasets = {"repe_honesty__plain_val", "ai_liar__original_without_answers"}
    violin_plot(
        exp,
        scorer,
        include_train=True,
        exclude_datasets=exclude_datasets,
        outfile=folder / "violin.png",
        figsize=(4.5, 2.5),
    )
    plot_roc_comparison(
        exp,
        scorer,
        exclude_datasets=exclude_datasets,
        outfile=folder / "roc_both.png",
        figsize=(4.5, 2.5),
        x_maxs=(1.0, 1.0),
        x_line=None,
    )


def followup_plots() -> None:
    exp = Experiment(ExperimentConfig.from_path(RESULTS_DIR / "followup"))
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01])
    folder = OUT_DIR / "probes/followup"
    folder.mkdir(parents=True, exist_ok=True)
    violin_plot(exp, scorer, outfile=folder / "violin.png", figsize=(4.5, 2.5))
    plot_roc_comparison(
        exp,
        scorer,
        outfile=folder / "roc_both.png",
        figsize=(4.5, 2.5),
        x_maxs=(1.0, 0.01),
        x_line=None,
    )
    ai_liar_violin(exp, scorer, outfile=folder / "ai_liar_violin.png", figsize=(4, 2.5), title="")


def sweep_repe() -> None:
    out_dir = OUT_DIR / "sweeps"
    out_dir.mkdir(parents=True, exist_ok=True)
    plot_reg_sweep(
        RESULTS_DIR / "repe",
        outfile=out_dir / "reg_coeff.png",
        y_mins=(0.94, 0.89),
        figsize=(7, 4),
    )


def sweep_sae_rp() -> None:
    out_dir = OUT_DIR / "sweeps"
    out_dir.mkdir(parents=True, exist_ok=True)
    plot_reg_sweep(
        RESULTS_DIR / "sae_rp",
        outfile=out_dir / "sae_rp_reg_coeff.png",
        y_mins=(0, 0),
        figsize=(7, 4),
    )


def comparison_plot():
    make_comparison_bar_chart(
        results_folder=RESULTS_DIR,
        label_key="name",
        # legend_title="Probe training set",
        value_names={
            "instructed_pairs": "Train on Instructed-Pairs",
            "roleplaying": "Train on Roleplaying",
            "sae_rp": "Train on Roleplaying:\nSAE activations",
            "followup": "Train on Roleplaying:\nFollow-up question",
            "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",
            "repe_honesty__plain_val",
            "ai_audit__reasoning",
            "ai_audit__a1_2_reasoning",
            "ai_audit__a1_2",
        },
        exp_for_black_box=RESULTS_DIR / "repe",
        y_lim=(0.0, 1.0),
        figsize=(6, 3),
        use_recall_for_control=True,
        outfile=OUT_DIR / "sweeps" / "trainset.png",
    )


def method_comparison_plot():
    make_comparison_bar_chart(
        results_folder=RESULTS_DIR / "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",
            "ai_audit__reasoning",
            "ai_audit__a1_2_reasoning",
            "ai_audit__a1_2",
        },
        y_lim=(0.0, 1.0),
        figsize=(7, 3.5),
        use_recall_for_control=True,
        outfile=OUT_DIR / "sweeps" / "method.png",
    )


def only_eot_comparison():
    make_comparison_bar_chart(
        results_folder=RESULTS_DIR / "vary_only_eot",
        label_key="name",
        value_names={
            "roleplaying": "entire response",
            "roleplaying_only_eot": "last token",
        },
        use_recall_for_control=True,
        outfile=OUT_DIR / "sweeps" / "last_token_only.png",
        figsize=(7, 3.5),
        separate_val_dataset_name=True,
    )


def reduction_comparison():
    make_reduction_comparison_bar_chart(
        results_folder=RESULTS_DIR / "repe",
        score_methods=["mean", "relu", "max"],
        outfile=OUT_DIR / "sweeps" / "reduction.png",
        figsize=(4, 3),
        excluded_datasets={
            "ai_liar__original_without_answers",
            "repe_honesty__plain_val",
            "ai_audit__no_reasoning",
            "ai_audit__reasoning",
            "ai_audit__a1_2_reasoning",
            "ai_audit__a1_2",
        },
        y_mins=(0, 0),
        legend=False,
    )
    make_reduction_comparison_bar_chart(
        results_folder=RESULTS_DIR / "roleplaying",
        score_methods=["mean", "relu", "max"],
        outfile=OUT_DIR / "sweeps" / "reduction-rp.png",
        figsize=(4.5, 3),
        excluded_datasets={
            "ai_liar__original_without_answers",
            "repe_honesty__plain_val",
            "ai_audit__no_reasoning",
            "ai_audit__reasoning",
            "ai_audit__a1_2_reasoning",
            "ai_audit__a1_2",
        },
        y_mins=(0, 0),
        relu_threshold=3.5,
    )


def grid_plot():
    """Apendix plot with detailed results for each probe."""
    out_dir = OUT_DIR / "probes"
    out_dir.mkdir(parents=True, exist_ok=True)

    probe_configs = [
        ProbeConfig(
            name="roleplaying",
            title="Train on Roleplaying",
            exp_path=RESULTS_DIR / "roleplaying",
        ),
        ProbeConfig(
            name="sae_rp",
            title="Train on Roleplaying: SAE activations",
            exp_path=RESULTS_DIR / "sae_rp",
        ),
        ProbeConfig(
            name="followup",
            title="Train on Roleplaying: Follow-up question",
            exp_path=RESULTS_DIR / "followup",
        ),
        ProbeConfig(
            name="descriptive",
            title="Train on single descriptive question",
            exp_path=RESULTS_DIR / "descriptive",
        ),
    ]

    exclude_datasets = {
        "repe_honesty__plain_val",
        "ai_liar__original_without_answers",
        "ai_audit__no_reasoning",
        "ai_audit__reasoning",
        "ai_audit__a1_2_reasoning",
        "ai_audit__a1_2",
    }

    plot_grid_of_detailed_results(
        probe_configs=probe_configs,
        exclude_datasets=exclude_datasets,
        output_path=OUT_DIR / "probes" / "comparison_detailed.png",
    )


def all_plots() -> None:
    repe_plots()
    grid_plot()
    comparison_plot()
    method_comparison_plot()
    only_eot_comparison()


if __name__ == "__main__":
    fire.Fire(
        {
            "all": all_plots,
            "repe": repe_plots,
            "roleplaying": roleplaying_plots,
            "sae_rp": sae_rp_plots,
            # "sae_repe": sae_repe_plots,
            "descriptive": descriptive_plots,
            "sweep_repe": sweep_repe,
            "sweep_sae_rp": sweep_sae_rp,
            "comparison": comparison_plot,
            "method": method_comparison_plot,
            "eot": only_eot_comparison,
            "reduction": reduction_comparison,
            "followup": followup_plots,
            "grid": grid_plot,
        }
    )
