#!/usr/bin/env python3
"""
LLM-as-judge for the Calm Coffee Co. brand voice fine-tune.

What it does:
  - Reads a CSV with columns: prompt, vanilla_output, finetuned_output
  - Sends each pair to Claude with blind A/B randomization
  - Claude scores both outputs 1-5 on four brand voice dimensions
  - Prints a win-rate table and saves full results to judge_results.json

Setup:
  pip install anthropic
  export ANTHROPIC_API_KEY=sk-ant-...

Run:
  python calm-coffee-judge.py calm-coffee-eval-template.csv
"""

import csv
import json
import os
import random
import sys
from collections import Counter

import anthropic

# The judge prompt — this is the rubric Claude scores against.
# Edit this if you change the brand voice rules.
JUDGE_PROMPT = """You are evaluating two product-copy outputs against the Calm Coffee Co. brand voice. Be a tough, honest judge — most marketing copy is generic, and we want to find which output is genuinely on-brand.

CALM COFFEE CO. BRAND VOICE RULES:
- All lowercase (except real proper nouns like origin names)
- Zero exclamation marks
- Sentence fragments dominate. Short. Periodic. Rhythmic.
- 4 lines or fewer
- Uses ritual or slowness language: "slow", "morning", "ritual", "quiet", "for the", "before", "sip", "no rush"
- Concrete sensory specifics: named flavor notes (citrus, cocoa, jasmine, stone fruit, brown sugar). Never generic adjectives like "exquisite," "premium," "delicious."
- Restrained. No marketing hype. No "introducing." No "experience the." Lets the cup speak.

THE PROMPT BOTH MODELS RECEIVED:
{prompt}

OUTPUT A:
{output_a}

OUTPUT B:
{output_b}

Score each output 1-5 on these four dimensions:
  on_brand_voice  - how strongly it follows the voice rules above
  specificity     - concrete sensory detail vs generic adjectives
  restraint       - does less, says less, no hype
  cadence         - reads aloud well; fragment rhythm

Then pick a winner: "A", "B", or "TIE". Explain in one sentence.

Respond ONLY with valid JSON in this exact shape (no markdown, no commentary):
{{
  "output_a": {{"on_brand_voice": <int 1-5>, "specificity": <int 1-5>, "restraint": <int 1-5>, "cadence": <int 1-5>}},
  "output_b": {{"on_brand_voice": <int 1-5>, "specificity": <int 1-5>, "restraint": <int 1-5>, "cadence": <int 1-5>}},
  "winner": "A",
  "reason": "<one sentence>"
}}"""


def judge_pair(client, prompt, vanilla_output, finetuned_output, judge_model="claude-sonnet-4-6"):
    """Blind-judge a single (vanilla, finetuned) pair."""
    # Randomize slot A vs slot B so the judge can't be biased by order
    if random.random() < 0.5:
        output_a, output_b = vanilla_output, finetuned_output
        a_is_finetuned = False
    else:
        output_a, output_b = finetuned_output, vanilla_output
        a_is_finetuned = True

    message = client.messages.create(
        model=judge_model,
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": JUDGE_PROMPT.format(
                prompt=prompt, output_a=output_a, output_b=output_b
            ),
        }],
    )

    raw = message.content[0].text.strip()
    # Strip markdown fencing if Claude added it
    if raw.startswith("```"):
        raw = raw.split("```")[1]
        if raw.startswith("json"):
            raw = raw[4:]
        raw = raw.strip("` \n")

    result = json.loads(raw)

    # Un-blind: map slot A/B back to vanilla vs finetuned
    if a_is_finetuned:
        finetuned_scores = result["output_a"]
        vanilla_scores = result["output_b"]
        winner_map = {"A": "finetuned", "B": "vanilla", "TIE": "tie"}
    else:
        vanilla_scores = result["output_a"]
        finetuned_scores = result["output_b"]
        winner_map = {"A": "vanilla", "B": "finetuned", "TIE": "tie"}

    return {
        "prompt": prompt,
        "vanilla_scores": vanilla_scores,
        "finetuned_scores": finetuned_scores,
        "winner": winner_map[result["winner"]],
        "reason": result["reason"],
    }


def main():
    if len(sys.argv) < 2:
        print("usage: python calm-coffee-judge.py outputs.csv")
        sys.exit(1)

    if "ANTHROPIC_API_KEY" not in os.environ:
        print("error: set ANTHROPIC_API_KEY first")
        print("  export ANTHROPIC_API_KEY=sk-ant-...")
        sys.exit(1)

    client = anthropic.Anthropic()
    results = []

    with open(sys.argv[1]) as f:
        reader = csv.DictReader(f)
        for i, row in enumerate(reader, 1):
            if not row.get("vanilla_output") or not row.get("finetuned_output"):
                print(f"[{i}] skipping (missing outputs): {row['prompt'][:60]}...")
                continue
            print(f"[{i}] judging: {row['prompt'][:60]}...")
            r = judge_pair(client, row["prompt"], row["vanilla_output"], row["finetuned_output"])
            results.append(r)

    if not results:
        print("\nno rows had both outputs filled in. fill the CSV and rerun.")
        sys.exit(1)

    n = len(results)
    winners = Counter(r["winner"] for r in results)

    def avg(side, dim):
        return sum(r[side][dim] for r in results) / n

    print("\n" + "=" * 60)
    print(f"RESULTS  ·  {n} prompts judged")
    print("=" * 60)

    print("\nWin rate:")
    print(f"  fine-tuned:  {winners['finetuned']}/{n}  ({100 * winners['finetuned'] / n:.0f}%)")
    print(f"  vanilla:     {winners['vanilla']}/{n}  ({100 * winners['vanilla'] / n:.0f}%)")
    print(f"  tie:         {winners['tie']}/{n}  ({100 * winners['tie'] / n:.0f}%)")

    print("\nAverage scores (1-5):")
    print(f"  {'dimension':<18} {'vanilla':>10} {'fine-tuned':>12}  {'delta':>8}")
    print(f"  {'-' * 18} {'-' * 10} {'-' * 12}  {'-' * 8}")
    for dim in ["on_brand_voice", "specificity", "restraint", "cadence"]:
        v = avg("vanilla_scores", dim)
        f_score = avg("finetuned_scores", dim)
        delta = f_score - v
        print(f"  {dim:<18} {v:>10.2f} {f_score:>12.2f}  {delta:>+8.2f}")

    avg_brand = avg("finetuned_scores", "on_brand_voice")
    win_rate = winners["finetuned"] / n
    passed = avg_brand >= 4.0 and win_rate >= 0.7
    print(f"\n{'PASS' if passed else 'FAIL'} — bar: fine-tuned on-brand avg ≥4.0 AND win rate ≥70%")
    print(f"  on-brand average: {avg_brand:.2f}  (need ≥4.0)")
    print(f"  fine-tuned win rate: {100 * win_rate:.0f}%  (need ≥70%)")

    with open("judge_results.json", "w") as f:
        json.dump(results, f, indent=2)
    print("\nfull results saved to judge_results.json")


if __name__ == "__main__":
    main()