# %% [markdown]
# # Evaluate on TruthfulQA
# Here we evaluate on TruthfulQA, to get a baseline with 1 "which is True" and 2. the learned label description
# %%
from src.data.truthfulqa import load_truthfulqa
from loguru import logger
import pandas as pd
from tqdm.auto import tqdm
import asyncio
import os
from dotenv import load_dotenv
load_dotenv()
from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices

import nest_asyncio
nest_asyncio.apply()

data = load_truthfulqa()

# HACK
# data = data[:32]


async def eval_tqa(instructions="Is the statement True or False? Answer A for True and B for False."):
    total_cost = 0.0

    res = []
    for row in tqdm(data):

        messages = [
            {"role": "user", "content": f"{instructions}\n\n{row['prompt']}\nAnswer?"},
            {"role": "assistant", "content": "Answer:"}
        ]

        response = await openrouter_completion_wlogprobs(
            model_id="qwen/qwen3-235b-a22b-2507",
            provider_whitelist=('Chutes','Nebius',),
            messages=messages,
            max_tokens=5,
            temperature=0.4,
            top_logprobs=8,
        )
        choice_strs = ["A", "B"]
        choice_logp, top_logp = get_logprobs_choices(response, choice_strs, lower=False)

        score = choice_logp["A"] - choice_logp["B"]
        res.append({
            "uid": row['uid'],
            "pred": 1 if score > 0 else 0,
            "label": row['vanilla_label'],
            "score": score,
            "choice_logp": choice_logp,
        })

        total_cost += response.get('usage', {}).get('cost', 0.0)

    return res, total_cost

# %%
# accuracy
res, total_cost = asyncio.run(eval_tqa())
print(f'Total cost ${total_cost:.2f}')

df_res = pd.DataFrame(res)
acc = (df_res['label'] == df_res['pred']).mean()
print(f"acc = {acc:2.2f}")

# %%


cleaned_instructions = """
1.  Set A – “Factually accurate / scientifically supported / myth-busting claims”  
2.  Set B – “Common misconceptions, falsehoods, myths, oversimplifications, or culturally based but inaccurate beliefs”
"""
res1, total_cost = asyncio.run(eval_tqa(cleaned_instructions))
print(f'Total cost ${total_cost:.6f}')


df_res1 = pd.DataFrame(res1)
df_res1
acc = (df_res1['label'] == df_res1['pred']).mean()
print(f"acc = {acc:2.2f}")
# acc = 0.83
# %%