# %%
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from deception_detection.rollouts import load_rollout_dialogues
from deception_detection.scripts.experiment import (
    Experiment,
    ExperimentConfig,
    create_prompt_scorer,
)
from deception_detection.types import Dialogue
from deception_detection.utils import REPO_ROOT, get_model_and_tokenizer

# %%
rollout_dialogues = load_rollout_dialogues(REPO_ROOT / "data/ft_fact_paired/out_simple_4_many.json")
experiment_path = REPO_ROOT / "results/lr__repe_honesty__gemma9b__20240910_151809"

dialogues: list[Dialogue] = []
honesties: list[str] = []
for dialogue in rollout_dialogues:
    for i, output_message in enumerate(dialogue.output_messages):
        dialogues.append(Dialogue(dialogue.input_messages + [output_message]))
        honesties.append(dialogue.model_graded_honesty[i])
rollout_dialogues = dialogues


cfg = ExperimentConfig.from_path(experiment_path)
_, tokenizer = get_model_and_tokenizer(cfg.type_model, cfg.size_model, omit_model=True)
experiment = Experiment(cfg, model=None, tokenizer=tokenizer)
scorer = create_prompt_scorer("mean", experiment, fpr=0.05)
scores = experiment.get_control_scores()["ft_fact_paired_rollouts"]
num = len(rollout_dialogues) // 2
honest_scores = scores.scores[:num]
deceptive_scores = scores.scores[num:]
honest_prompt_scores = scorer.prompt_scores(honest_scores)
deceptive_prompt_scores = scorer.prompt_scores(deceptive_scores)
honest_honesties = honesties[: len(honest_prompt_scores)]
deceptive_honesties = honesties[len(honest_prompt_scores) :]

print(len(honest_prompt_scores))
print(len(deceptive_prompt_scores))
print(len(rollout_dialogues))

honest_rollout_dialogues = rollout_dialogues[: len(honest_prompt_scores)]
deceptive_rollout_dialogues = rollout_dialogues[len(honest_prompt_scores) :]

df = pd.DataFrame(columns=["role", "honest", "score"])  # type: ignore

for i, _dialogue in enumerate(honest_rollout_dialogues):
    true = honest_honesties[i] == "honest"
    df.loc[len(df)] = {  # type: ignore
        "role": "alice",
        "honest": true,
        "score": honest_prompt_scores[i].item(),
    }
for i, _dialogue in enumerate(deceptive_rollout_dialogues):
    true = deceptive_honesties[i] == "honest"
    df.loc[len(df)] = {  # type: ignore
        "role": "bob",
        "honest": true,
        "score": deceptive_prompt_scores[i].item(),
    }


# %%

# Create a figure with subplots
fig = make_subplots(rows=1, cols=1)

# Define colors for honest and dishonest
colors = {"True": "green", "False": "red"}


def get_name(role: str, honest: bool) -> str:
    counts = df.groupby(["role", "honest"]).size().to_dict()
    if role == "alice" and honest:
        return f"Alice (honest) (n={counts[(role, honest)]})"
    elif role == "alice" and not honest:
        return f"Alice (deceptive) (n={counts[(role, honest)]})"
    elif role == "bob" and honest:
        return f"Bob (honest) (n={counts[(role, honest)]})"
    elif role == "bob" and not honest:
        return f"Bob (deceptive) (n={counts[(role, honest)]})"
    else:
        return "Unknown"


# Create a single violin plot for all data
fig.add_trace(
    go.Violin(
        x=df.apply(lambda row: get_name(row["role"], row["honest"]), axis=1),  # type: ignore
        y=df["score"],
        name="All Data",
        box_visible=True,
        meanline_visible=True,
        opacity=0.6,
        line_color="black",
    )
)

# Update layout
fig.update_layout(
    title="Distribution of Scores by Role and Honesty",
    xaxis_title="Role and Honesty",
    yaxis_title="Score",
    violinmode="group",
    violingap=0.1,
    violingroupgap=0.05,
)

# Show the plot
fig.show()

# %%
