#!/usr/bin/env python3
"""Quick 3×3 test — NIMA + OCR metrics comparison.

3 best models × 3 prompts (1 portrait, 1 complex scene, 1 text rendering).
Tests new VLM 0.1 increment scoring + NIMA + OCR metrics.

Expected cost: ~$0.10-0.15
"""

import os
import sys
import time
from pathlib import Path

# Load .env
env_file = Path(__file__).resolve().parent.parent.parent / ".env"
if env_file.exists():
    for line in env_file.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            k, v = line.split("=", 1)
            os.environ.setdefault(k.strip(), v.strip())

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from evalytic.bench.runner import run_bench
from evalytic.report.terminal import print_full_report
from evalytic.report.json_report import write_json

# ---------- Config ----------

# Best 3 models from previous leaderboard run
MODELS = [
    "ideogram-v3",        # Ideogram, $0.080, top text rendering
    "recraft-v3",         # Recraft, $0.040, strong overall
    "flux-pro",           # BFL, $0.040, high quality
]

# 3 diverse prompts: portrait (no text), complex scene (no text), text rendering (has text)
PROMPTS = [
    {
        "item_id": "test-port",
        "prompt": "A bearded blacksmith in a dark workshop, sparks flying from the anvil, cinematic side lighting",
        "has_text": False,
    },
    {
        "item_id": "test-scene",
        "prompt": "A vintage coffee shop storefront with hand-lettered chalkboard menu, morning sunlight, 35mm film photography style",
        "has_text": False,
    },
    {
        "item_id": "test-text",
        "prompt": "A glowing neon sign reading OPEN 24 HOURS in a dark rainy alley, reflections on wet cobblestones",
        "has_text": True,
        "expected_text": "OPEN 24 HOURS",
    },
]

# ---------- Pre-flight ----------

print("=== NIMA + OCR Test — 3 prompts × 3 models ===")
print(f"Models: {', '.join(MODELS)}")
print(f"Prompts: {len(PROMPTS)}")
print(f"Metrics: clip + sharpness + nima + ocr")
print(f"Judge: Consensus (gemini-2.5-flash + gpt-5.2)")
print(f"VLM scoring: 0.1 increments (NEW)")
print(f"FAL_KEY: {'set' if os.environ.get('FAL_KEY') else 'MISSING!'}")
print(f"GEMINI_API_KEY: {'set' if os.environ.get('GEMINI_API_KEY') else 'MISSING!'}")
print()

# ---------- Run ----------

start = time.monotonic()

report = run_bench(
    models=MODELS,
    prompts=PROMPTS,
    seed=42,
    judges=["gemini-2.5-flash", "gpt-5.2"],
    metrics=["clip", "sharpness", "nima", "ocr"],
    concurrency=6,
    timeout=600,
    name="nima-ocr-test-3x3",
    verbose=True,
)

elapsed = time.monotonic() - start

# ---------- Terminal output ----------

print_full_report(report)

# ---------- Detailed metric comparison ----------

print("\n" + "=" * 80)
print("DETAILED METRIC COMPARISON")
print("=" * 80)

for item in report.items:
    print(f"\n--- Prompt: {item.prompt[:80]}... ---")
    for model, ir in sorted(item.results.items()):
        if ir.status != "success":
            print(f"  {model}: FAILED")
            continue

        # VLM scores
        vlm_str = ", ".join(f"{s.dimension}={s.score}" for s in ir.scores)

        # Metrics
        metrics_str = ", ".join(f"{m.metric}={m.value:.4f}" for m in ir.metrics)

        print(f"  {model}:")
        print(f"    VLM: {vlm_str}")
        print(f"    Metrics: {metrics_str}")

# ---------- Summary table ----------

print("\n" + "=" * 80)
print("SUMMARY TABLE")
print("=" * 80)
print(f"{'Model':<25} {'VQ':>5} {'PA':>5} {'TR':>5} {'CLIP':>7} {'Sharp':>7} {'NIMA':>7} {'OCR':>7} {'Overall':>8}")
print("-" * 80)

for model, ms in sorted(report.summary.items(), key=lambda x: x[1].overall_score, reverse=True):
    vq = ms.dimension_averages.get("visual_quality", 0)
    pa = ms.dimension_averages.get("prompt_adherence", 0)
    tr = ms.dimension_averages.get("text_rendering", 0)
    clip = ms.metric_averages.get("clip_score", 0)
    sharp = ms.metric_averages.get("sharpness", 0)
    nima = ms.metric_averages.get("nima_score", 0)
    ocr = ms.metric_averages.get("ocr_accuracy", 0)
    print(f"{model:<25} {vq:>5.2f} {pa:>5.2f} {tr:>5.2f} {clip:>7.4f} {sharp:>7.4f} {nima:>7.4f} {ocr:>7.4f} {ms.overall_score:>8.3f}")

# ---------- Output files ----------

out_dir = Path("scripts/output")
out_dir.mkdir(parents=True, exist_ok=True)

json_path = out_dir / "nima-ocr-test-3x3.json"
write_json(report, str(json_path))
print(f"\nJSON: {json_path.resolve()}")

# ---------- Summary ----------

total_cost = report.cost.generation_total_usd + report.cost.judge_total_usd
print(f"\nTotal: {elapsed:.0f}s, ${total_cost:.3f}")
