# %%
from termcolor import colored

from deception_detection.data import WerewolfOpenResponse
from deception_detection.rollouts import load_rollout_dialogues
from deception_detection.utils import REPO_ROOT

werewolf_rollouts_path = REPO_ROOT / "werewolf_reasoning"
werewolf_rollouts = load_rollout_dialogues(werewolf_rollouts_path)

werewolf_dataset = WerewolfOpenResponse()


for i, rollout in enumerate(werewolf_rollouts):
    role = werewolf_dataset.roles[i]
    color = "red" if role in ["Werewolf", "Minion"] else "green"
    print("Role:", colored(role, color))
    print("Night Info:", werewolf_dataset.night_infos[i])
    print("Transcript:")
    for message in werewolf_dataset.discussions[i]:
        user, text = next(message.items().__iter__())
        print(f"  {user}: {text}")

    print("\nCompletions:")
    for completion, grading in zip(rollout.output_strs, rollout.model_graded_honesty, strict=True):
        color = "green" if grading == "honest" else "red"
        print(colored(f" - {completion}", color))
    assert rollout.honesty_percentage is not None
    print(f"Honesty percentage: {rollout.honesty_percentage*100:.2f}%")
    print("\n\n")

# %%
