"""
This script evaluates the performance of a detector at each layer of a model.

Will result in an experiment folder, although it doesn't compute the full results.

Some parameters are hardcoded in this script (in particular use of llama-70b)
"""

# %%
import json
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import torch
from dotenv import load_dotenv
from fire import Fire
from tqdm import tqdm, trange

from deception_detection.activations import PairedActivations
from deception_detection.data import DialogueDataset
from deception_detection.detectors import LogisticRegressionDetector
from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.log import logger
from deception_detection.metrics import get_auroc
from deception_detection.models import ModelName
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.results.plot import save_plot_exp
from deception_detection.results.utils import get_dataset_color, to_readable_dataset_name
from deception_detection.scores import Scores
from deception_detection.tokenized_data import TokenizedDataset
from deception_detection.types import Label

load_dotenv()


# %%
def get_experiment(folder: str | Path):
    cfg = ExperimentConfig.from_path(Path(folder))
    return Experiment(cfg)


def get_detectors_for_different_reg_coeffs(
    experiment: Experiment, reg_coeffs: list[float]
) -> list[LogisticRegressionDetector]:
    experiment._load_model()  # type: ignore
    assert experiment.model is not None
    split_paired_acts = PairedActivations.from_dataset(
        dataset=experiment.train_dataset,
        tokenizer=experiment.tokenizer,
        model=experiment.model,
        detect_only_start_of_turn=experiment.cfg.detect_only_start_of_turn,
        detect_only_last_token=experiment.cfg.detect_only_last_token,
        layers=experiment.cfg.detect_layers,
        use_goodfire_sae_acts=experiment.cfg.use_goodfire_sae_acts,
    )
    honest_fit_acts, deceptive_fit_acts = split_paired_acts.get_all_acts()
    detectors = []
    for c in tqdm(reg_coeffs, desc="Fitting detectors"):
        detector = LogisticRegressionDetector(
            layers=experiment.cfg.detect_layers,
            reg_coeff=c,
            normalize=experiment.cfg.normalize_acts,
        )
        detector.fit(deceptive_fit_acts, honest_fit_acts)
        detectors.append(detector)
    return detectors


def get_scores_shard(
    experiment: Experiment, dataset: DialogueDataset, detectors: list[LogisticRegressionDetector]
) -> list[Scores]:
    toks = TokenizedDataset.from_dataset(dataset, experiment.tokenizer)
    assert experiment.model is not None
    acts = experiment.get_acts(toks)
    scores_by_detector = [detector.score(acts) for detector in detectors]
    return scores_by_detector


def get_scores(
    experiment: Experiment,
    dataset: DialogueDataset,
    detectors: list[LogisticRegressionDetector],
    shard_size: int = 20,
) -> dict[float, Scores]:
    logger.info(f"Getting scores for {dataset.name} with shard size {shard_size}")
    scores_by_shard_by_layer: list[list[Scores]] = []
    if experiment.cfg.use_goodfire_sae_acts:
        shard_size = len(dataset)
    for start_idx in trange(0, len(dataset), shard_size, desc="Dataset Shards"):
        shard_dataset = dataset[start_idx : start_idx + shard_size]
        scores_by_shard_by_layer.append(get_scores_shard(experiment, shard_dataset, detectors))

    scores_by_detector = [
        Scores.from_list(list(scores)) for scores in zip(*scores_by_shard_by_layer, strict=True)
    ]
    for scores in scores_by_detector:
        scores.labels = dataset.labels

    return {
        detector.reg_coeff: scores
        for detector, scores in zip(detectors, scores_by_detector, strict=True)
    }


def compute_and_save_scores(
    experiment: Experiment,
    dataset: DialogueDataset,
    detectors: list[LogisticRegressionDetector],
    shard_size: int = 20,
):
    path = experiment.cfg.folder / "scores_by_coeff.json"
    path.parent.mkdir(exist_ok=True)
    data: dict[str, dict[float, Scores | Any]]
    if path.exists():
        with open(path) as f:
            data = json.load(f)
    else:
        data = {}
    if dataset.name not in data:
        data[dataset.name] = get_scores(experiment, dataset, detectors, shard_size)
    with open(path, "w") as f:
        json.dump(data, f, default=lambda x: x.to_dict())


def run(exp_path: str | Path | None = None):
    if exp_path is None:
        config = ExperimentConfig(
            method="lr",
            model_name=ModelName.LLAMA_70B_3_3,
            train_data="repe_honesty__plain",
            eval_data=[],
            eval_on_policy=True,
            control_data=["alpaca__plain"],
            control_dataset_size=1000,
            id="by_reg_coeff",
            detect_layers=[22],
        )
        experiment = Experiment(config)
    else:
        experiment = get_experiment(exp_path)

    coeffs: list[float] = (10 ** torch.arange(-2, 5.1, 0.5)).tolist()
    detectors = get_detectors_for_different_reg_coeffs(experiment, coeffs)
    excluded_datasets = {"ai_liar__original_without_answers"}
    for dataset in experiment.eval_datasets:
        if dataset.name in excluded_datasets:
            continue
        shard_size = {
            "repe_honesty__plain": 20,
            "repe_honesty__plain_val": 20,
            "roleplaying__plain": 20,
            "roleplaying__plain_val": 20,
            "harm_refusal_choice__plain": 20,
            "harm_refusal_choice__plain_val": 20,
            "harm_refusal_knowledge_report__plain": 20,
            "harm_refusal_knowledge_report__plain_val": 20,
            "insider_trading_XYZ__plain": 20,
            "insider_trading_XYZ__plain_val": 20,
            "instruct_dishonesty_default__plain": 20,
            "instruct_dishonesty_default__plain_val": 20,
            "instruct_dishonesty_donotadmit__plain": 20,
            "instruct_dishonesty_donotadmit__plain_val": 20,
            "mask__plain": 20,
            "mask__plain_val": 20,
            "sleeper_agents__plain": 20,
            "sleeper_agents__plain_val": 20,
            "convince_me__plain": 20,
            "convince_me__plain_val": 20,
            "insider_trading__upscale": 4,
            "insider_trading_doubledown__upscale": 4,
            "sandbagging_v2__wmdp_mmlu": 2,
            "alpaca__plain": 4,
        }[dataset.name]
        compute_and_save_scores(experiment, dataset, detectors, shard_size=shard_size)
    for dataset in experiment.control_datasets:
        shard_size = 4
        compute_and_save_scores(experiment, dataset, detectors, shard_size=shard_size)
    plot(experiment.cfg.folder)


def plot(
    exp_path: str | Path,
    use_recall: bool = True,
    outfile: str | Path | None = None,
    y_mins: tuple[float, float] = (0.8, 0.65),
    figsize: tuple[float, float] = (7, 4),
):
    # Load the scores data
    with open(Path(exp_path) / "scores_by_coeff.json") as f:
        data = json.load(f)

    # Get AUROCs for each layer
    prompt_scorer = create_prompt_scorer(method="mean", experiment=None, fprs=[])

    eval_scores_by_coeff: dict[str, dict[float, Scores]] = {
        dataset: {
            float(coeff): Scores.from_dict(scores) for coeff, scores in dataset_scores.items()
        }
        for dataset, dataset_scores in data.items()
        if (dataset != "alpaca__plain" and not dataset.endswith("_val"))
    }
    alpaca_scores_by_coeff: dict[float, Scores] = {
        float(coeff): Scores.from_dict(scores) for coeff, scores in data["alpaca__plain"].items()
    }

    def get_alpaca_auroc(scores: Scores, coeff: float) -> float:
        deceptive_scores = scores.get_scores_by_label()[Label.DECEPTIVE]
        return get_auroc(
            prompt_scorer.process(alpaca_scores_by_coeff[coeff], normalize=False),
            prompt_scorer.process(deceptive_scores, normalize=False),
        )

    def get_alpaca_recall(scores: Scores, coeff: float) -> float:
        deceptive_scores_by_prompt = prompt_scorer.process(
            scores.get_scores_by_label()[Label.DECEPTIVE], normalize=False
        )
        alpaca_scores_by_prompt = prompt_scorer.process(
            alpaca_scores_by_coeff[coeff], normalize=False
        )
        alpaca_threshold = alpaca_scores_by_prompt.quantile(0.99)
        return (deceptive_scores_by_prompt > alpaca_threshold).float().mean().item()

    _, axs = plt.subplots(2, 1, figsize=figsize, sharex=True)
    for dataset_name, scores_data in eval_scores_by_coeff.items():
        label = to_readable_dataset_name(dataset_name)
        color = get_dataset_color(dataset_name)
        aurocs = {coeff: prompt_scorer.auroc(scores) for coeff, scores in scores_data.items()}
        aurocs = dict(sorted(aurocs.items()))
        axs[0].semilogx(
            list(aurocs.keys()),
            list(aurocs.values()),
            marker="o",
            color=color,
            label=label,
        )
        if use_recall:
            alpaca_points = {
                coeff: get_alpaca_recall(scores, coeff) for coeff, scores in scores_data.items()
            }
        else:
            alpaca_points = {
                coeff: get_alpaca_auroc(scores, coeff) for coeff, scores in scores_data.items()
            }
        alpaca_points = dict(sorted(alpaca_points.items()))

        axs[1].semilogx(
            list(alpaca_points.keys()),
            list(alpaca_points.values()),
            color=color,
            marker="o",
        )

    axs[0].set_title("AUROC Deceptive vs Honest")
    axs[0].set_ylabel("AUROC")
    axs[0].set_ylim(y_mins[0], 1)
    axs[1].set_ylim(y_mins[1], 1)
    axs[0].grid(True)
    axs[0].axvline(x=10, color="k", linestyle="--", linewidth=1)

    if use_recall:
        axs[1].set_title(r"Recall @ 1% FRP on Control")
        axs[1].set_ylabel("Recall (TPR)")
    else:
        axs[1].set_title(r"AUROC Deceptive vs Control")
        axs[1].set_ylabel("AUROC")
    axs[1].set_xlabel(r"Regularization Coefficient $\lambda$")
    axs[1].grid(True)
    axs[1].axvline(x=10, color="k", linestyle="--", linewidth=1)

    axs[0].legend(bbox_to_anchor=(1.025, 1), loc="upper left")

    # Save the plot
    exp = Experiment(ExperimentConfig.from_path(Path(exp_path)))
    save_plot_exp(exp, "by_reg_coeff_auroc.png", Path(outfile) if outfile else None)


if __name__ == "__main__":
    Fire({"run": run, "plot": plot})
