# %%
from deception_detection.data import AILiarDataset, InsiderTrading, WerewolfOpenResponse
from deception_detection.rollouts import (
    RolloutGenerator,
    evaluate_all_dialogues,
    load_rollout_dialogues,
)
from deception_detection.utils import get_gemma_model_and_tokenizer

model, tokenizer = get_gemma_model_and_tokenizer(9)
assert model is not None

# %%
dataset = AILiarDataset()
rollout_dataset = RolloutGenerator.from_paired_dialogue_dataset(dataset, model, tokenizer)
rollout_dataset.generate_many(5)
rollout_dataset.grade_all_dialogues()
rollout_dataset.save("data/how_to_catch_an_ai_liar/rollouts.json")
rollout_dialogues = load_rollout_dialogues("data/how_to_catch_an_ai_liar/rollouts.json")
evaluate_all_dialogues(rollout_dialogues)

# %%
dataset = InsiderTrading()
rollout_dataset = RolloutGenerator.from_dialogue_dataset(dataset, model, tokenizer)
rollout_dataset.generate_many(5)
rollout_dataset.grade_all_dialogues(verbose=False)
rollout_dataset.save("../../data/insider_trading/rollouts5_27b.json")
rollout_dialogues = load_rollout_dialogues("../../data/insider_trading/rollouts5_27b.json")
evaluate_all_dialogues(rollout_dialogues)

# %%
dataset = WerewolfOpenResponse()
rollout_dataset = RolloutGenerator.from_dialogue_dataset(dataset, model, tokenizer)
rollout_dataset.generate_many(1)

# %%
