import openai
import anthropic
import pandas as pd
from tqdm import tqdm
import time

api_keys = {
    "openai": "OPENAI_KEY",
    "anthropic": "ANTHROPIC_KEY"
}

def openai_judge(prompt, answer, question, model='gpt-4o'):
    completion = openai.ChatCompletion.create(
        api_key=api_keys['openai'],
        model=model,
        temperature=0,
        messages=[
            {"role": "system", "content": prompt},
            {"role": "user", "content": f"Q: {question}\nA: {answer}"}
        ]
    )
    return completion['choices'][0]['message']['content'].strip()

def anthropic_judge(prompt, answer, question, model='claude-3-opus-20240229'):
    client = anthropic.Anthropic(api_key=api_keys['anthropic'])
    response = client.messages.create(
        model=model,
        max_tokens=256,
        messages=[
            {"role": "user", "content": prompt + f"\nQ: {question}\nA: {answer}"}
        ]
    )
    return response.content[0].text.strip()

def replay_evaluations(evals, judge_fn, runs=3, delay=2):
    results = []
    for idx, row in tqdm(evals.iterrows(), total=len(evals)):
        verdicts = []
        for r in range(runs):
            try:
                verdicts.append(judge_fn(
                    prompt=row['prompt'],
                    answer=row['agent_output'],
                    question=row['task']
                ))
                time.sleep(delay)
            except Exception:
                verdicts.append('ERROR')
        results.append(verdicts)
    return pd.DataFrame(results, columns=[f'run_{i+1}' for i in range(runs)])

# Load logs: [{"task": "...", "agent_output": "...", "prompt": "..."}]
evals = pd.read_json('agent_judgement_inputs.json')

# OpenAI judge replay
verdicts_openai = replay_evaluations(evals, openai_judge, runs=5, delay=2)
verdicts_openai.to_csv('openai_verdicts.csv')

# Anthropic judge replay
verdicts_anthropic = replay_evaluations(evals, anthropic_judge, runs=5, delay=5)
verdicts_anthropic.to_csv('anthropic_verdicts.csv')