import json
import shutil
from pathlib import Path
from typing import cast

import fire
import pandas as pd
import plotly.graph_objects as go
import seaborn as sns

from deception_detection.data.sandbagging_v2 import SandbaggingV2Dataset
from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.repository import DatasetRepository
from deception_detection.results.utils import get_dataframe_of_scores_on_dataset
from deception_detection.scores import Scores
from deception_detection.utils import REPO_ROOT


def sandbagging_data_labels_df() -> pd.DataFrame:
    """
    Create a DataFrame analyzing sandbagging behavior in the dataset responses.

    Args:
        dataset: A SandbaggingV2Dataset instance containing dialogues with responses

    Returns:
        pd.DataFrame: DataFrame with analysis of sandbagging behavior, containing columns:
            - id: message ID
            - step1_results: sandbag/normal decision
            - step2_results: initial answer
            - answer: final answer
            - did_sandbag: bool indicating if sandbagging was attempted
            - step2_answer: initial answer (alias)
            - final_answer: final answer (alias)
            - answer_equal: bool indicating if initial and final answers match
            - dataset: 'wmdp' or 'mmlu'
    """
    repo = DatasetRepository()
    dataset = cast(
        SandbaggingV2Dataset,
        repo.get("sandbagging_v2__wmdp_mmlu", "llama-70b-3.3", trim_reasoning=False),
    )

    messages = [d[-1].content for d in dataset.dialogues]
    extract_tag = dataset.extract_tag_content  # Use the class method

    def get_did_sandbag(m: str) -> bool | None:
        match extract_tag(m, "step1-results"):
            case "sandbag":
                return True
            case "normal":
                return False
            case _:
                return None

    assert dataset.metadata is not None
    return pd.DataFrame(
        {
            "id": dataset.metadata["id"],
            "step1_results": [extract_tag(m, "step1-results") for m in messages],
            "step2_results": [extract_tag(m, "step2-results") for m in messages],
            "answer": [extract_tag(m, "answer") for m in messages],
            "did_sandbag": [get_did_sandbag(m) for m in messages],
            "step2_answer": [extract_tag(m, "step2-results") for m in messages],
            "final_answer": [extract_tag(m, "answer") for m in messages],
            "answer_equal": [
                extract_tag(m, "step2-results") == extract_tag(m, "answer") for m in messages
            ],
            "dataset": [
                "wmdp" if expected_label == "deceptive" else "mmlu"
                for expected_label in dataset.metadata["expected_label"]
            ],
        }
    )


def make_sankey(df: pd.DataFrame) -> None:
    """
    Make sankey counting occurances between:
    - the "dataset" column, either "wmdp" or "mmlu"
    - the "did_sandbag" column, either "True", "False", or "None"
    - If did sandbag, the "answer_equal" column, either "True", "False"

    """

    df = sandbagging_data_labels_df()

    # Create node labels
    nodes = [
        "WMDP",
        "MMLU",  # Dataset nodes (0, 1)
        "Sandbag",
        "No Sandbag",
        "Invalid",  # Decision nodes (2, 3, 4)
        "Same Answer",
        "Different Answer",
    ]  # Outcome nodes (5, 6)

    # Calculate flows
    flows: list[tuple[int, int]] = []
    values: list[int] = []

    # Dataset -> Decision flows
    for dataset, did_sandbag in [
        ("wmdp", True),
        ("wmdp", False),
        ("wmdp", None),
        ("mmlu", True),
        ("mmlu", False),
        ("mmlu", None),
    ]:
        count = len(df[(df["dataset"] == dataset) & (df["did_sandbag"] == did_sandbag)])  # type: ignore
        if count > 0:
            source = 0 if dataset == "wmdp" else 1
            target = 2 if did_sandbag else (3 if did_sandbag is False else 4)
            flows.append((source, target))
            values.append(count)

    # Decision -> Outcome flows (only for sandbagging cases)
    sandbag_subset = df[df["did_sandbag"] == True]

    # Same answer flow
    count = len(sandbag_subset[sandbag_subset["answer_equal"] == True])  # type: ignore
    if count > 0:
        flows.append((2, 5))
        values.append(count)

    # Different answer flow
    count = len(sandbag_subset[sandbag_subset["answer_equal"] == False])  # type: ignore
    if count > 0:
        flows.append((2, 6))
        values.append(count)

    # Create the Sankey diagram
    fig = go.Figure(
        data=[
            go.Sankey(
                node=dict(
                    pad=15,
                    thickness=20,
                    line=dict(color="black", width=0.5),
                    label=nodes,
                    color=["grey", "grey", "grey", "blue", "grey", "red", "yellow"],
                ),
                link=dict(source=[f[0] for f in flows], target=[f[1] for f in flows], value=values),
            )
        ]
    )

    fig.update_layout(title_text="Sandbagging Analysis Flow", font_size=12, height=600)

    fig.show()


def plot_sandbagging_analysis() -> None:
    """
    Violin plot of scores for sandbagging and non-sandbagging dialogues.
    """
    exp_path = REPO_ROOT / "results/33/repe"
    exp = Experiment(ExperimentConfig.from_path(exp_path))
    scorer = create_prompt_scorer("mean", exp, fprs=[0.01])
    scores_df = get_dataframe_of_scores_on_dataset(exp, "sandbagging_v2__wmdp_mmlu", scorer)
    labels_df = sandbagging_data_labels_df()

    # merge on id column, use that as index
    df = scores_df.merge(labels_df, on="id", how="inner")

    def get_label(row: pd.Series) -> str:  # type: ignore
        if row["did_sandbag"]:  # type: ignore
            if row["answer_equal"]:  # type: ignore
                return "ambigious"
            else:
                return "sandbag"
        else:
            return "normal"

    df["label"] = df.apply(get_label, axis=1)  # type: ignore
    sns.violinplot(x="label", y="score", data=df, hue="dataset")


def update_labels_in_rollout_files(model_name: str = "llama-70b-3.3") -> None:
    """
    Update the labels in the rollout files using the `SandbaggingV2Dataset.label_dialogue` method.
    """
    repo = DatasetRepository()
    path = repo.dataset_paths[f"sandbagging_v2__wmdp_mmlu__{model_name}"]

    with open(path) as f:
        rollout_data = json.load(f)

    reasoning = "Computed from `SandbaggingV2Dataset.label_dialogue`"
    for rollout_item in rollout_data["rollouts"]:
        labels = [
            SandbaggingV2Dataset.label_response(message) for message in rollout_item["output_strs"]
        ]
        rollout_item["graded_deception"] = [label.value for label in labels]
        rollout_item["grader_reasoning"] = [reasoning] * len(labels)

    with open(path, "w") as f:
        json.dump(rollout_data, f, indent=2)

    print(f"Updated {path}")


def update_labels_in_experiment_scores(exp_path: str | Path = "all") -> None:
    """
    Update the labels in the experiment scores for all models.
    """
    if exp_path == "all":
        all_exp_paths = [
            REPO_ROOT / "results/33/repe",
            REPO_ROOT / "results/33/sae_rp",
            REPO_ROOT / "results/33/roleplaying",
            REPO_ROOT / "results/33/descriptive",
            REPO_ROOT / "results/by_method/cmms",
            REPO_ROOT / "results/by_method/lr",
            REPO_ROOT / "results/by_method/lat",
            REPO_ROOT / "results/by_method/mms",
            REPO_ROOT / "results/vary_only_eot/roleplaying",
            REPO_ROOT / "results/vary_only_eot/roleplaying_only_eot",
            REPO_ROOT / "results/by_layer",
        ]
        for exp_path in all_exp_paths:
            update_labels_in_experiment_scores(exp_path)
        return

    exp = Experiment(ExperimentConfig.from_path(Path(exp_path)))
    sandbagging_dataset = [d for d in exp.eval_datasets if d.name == "sandbagging_v2__wmdp_mmlu"][0]

    for score_file in ["scores.json", "scores_all_tokens.json", "black_box_scores.json"]:
        scores_path = exp.cfg.folder / score_file
        if not scores_path.exists():
            continue
        # backup the scores file if backup doesn't exist
        if not scores_path.with_suffix(".bak").exists():
            shutil.copy(scores_path, scores_path.with_suffix(".bak"))

        try:
            # update the scores file
            with open(scores_path) as f:
                scores_data = json.load(f)

            sandbagging_scores = Scores.from_dict(scores_data["sandbagging_v2__wmdp_mmlu"])
            assert len(sandbagging_scores.labels) == len(sandbagging_dataset.labels)
            sandbagging_scores.labels = sandbagging_dataset.labels
            scores_data["sandbagging_v2__wmdp_mmlu"] = sandbagging_scores.to_dict()

            with open(scores_path, "w") as f:
                json.dump(scores_data, f, indent=2)

            print(f"Updated {scores_path}")
        except Exception as e:
            # revert the scores file
            shutil.move(scores_path.with_suffix(".bak"), scores_path)
            raise e

    for score_file in ["scores_by_layer.json", "scores_by_coeff.json"]:
        scores_path = exp.cfg.folder / score_file
        if not scores_path.exists():
            continue

        if not scores_path.with_suffix(".bak").exists():
            shutil.copy(scores_path, scores_path.with_suffix(".bak"))

        try:
            with open(scores_path) as f:
                scores_data = json.load(f)

            for val, sb_score_data in scores_data["sandbagging_v2__wmdp_mmlu"].items():
                score_object = Scores.from_dict(sb_score_data)
                assert len(score_object.labels) == len(sandbagging_dataset.labels)
                score_object.labels = sandbagging_dataset.labels
                scores_data["sandbagging_v2__wmdp_mmlu"][val] = score_object.to_dict()

            with open(scores_path, "w") as f:
                json.dump(scores_data, f, indent=2)

            print(f"Updated {scores_path}")

        except Exception as e:
            # revert the scores file
            shutil.move(scores_path.with_suffix(".bak"), scores_path)
            raise e


if __name__ == "__main__":
    fire.Fire(
        {
            "violin": plot_sandbagging_analysis,
            "sankey": make_sankey,
            "update_rollout": update_labels_in_rollout_files,
            "update_scores": update_labels_in_experiment_scores,
        }
    )
