# finetune/eval_ft.py
"""Stage 5 gate — side-by-side Bengali eval: fine-tuned MiniCPM-V (vLLM) vs the
native base model (the current shipping "Lever #1" path).

CLAUDE.md / finetune/README are explicit: ship the fine-tune ONLY if it clearly
beats the native path in a human read. The train metrics (loss/token_acc) say the
LoRA fit the distilled targets — they do NOT say Bengali quality improved. This
script produces the artifact a Bengali speaker needs to make that call before
FINETUNED_VISION_MODEL is ever set.

Fairness: BOTH paths get the EXACT same app-built Bengali prompt (build_story_prompt)
and the same image — exactly what each would receive in production. The only
difference under test is the model weights.

  native: core.modal_infra.generate_story_remote  → base openbmb/MiniCPM-V-4_5 (Ollama)
  FT:     finetune.serve_vllm.generate_story_ft_remote → merged LoRA (vLLM)

Held-out set: the 61 labelset images that the purity gate rejected, so they were
NEVER trained on, yet are on-distribution. (Override with --images for your own.)

Run:
    uv run modal deploy finetune/serve_vllm.py          # FT server must be live
    uv run python finetune/eval_ft.py --n 10            # 10 held-out images
    uv run python finetune/eval_ft.py --images path/to/dir --n 8 --style রূপকথা
Out:
    finetune/eval_results/ft_vs_native_YYYYMMDD_HHMM.md
"""

import argparse
import base64
import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.model_config import ACTIVE_STACK, get_vision_options
from core.modal_infra import generate_story_remote
from core.prompts import STYLES, build_story_prompt
from finetune.serve_vllm import generate_story_ft_remote

TRAIN_JSON = Path("finetune/data/train.json")
LABELSET = Path("finetune/data/labelset")


def held_out_images() -> list[Path]:
    """Labelset images that are NOT in train.json — unseen but on-distribution."""
    trained = {x["image"].split("/")[-1] for x in json.loads(TRAIN_JSON.read_text())}
    imgs = [
        p for p in sorted(LABELSET.glob("*"))
        if p.suffix.lower() in (".jpg", ".jpeg", ".png") and p.name not in trained
    ]
    return imgs


def _encode(path: Path) -> str:
    return base64.b64encode(path.read_bytes()).decode()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--images", default=None, help="dir of images (default: held-out labelset)")
    ap.add_argument("--n", type=int, default=10, help="number of images to evaluate")
    ap.add_argument("--style", default="রূপকথা", choices=list(STYLES["bn"].keys()))
    ap.add_argument("--instruction", default="একটা গল্প বলো")
    args = ap.parse_args()

    if args.images:
        imgs = [
            p for p in sorted(Path(args.images).glob("*"))
            if p.suffix.lower() in (".jpg", ".jpeg", ".png")
        ]
    else:
        imgs = held_out_images()
    imgs = imgs[: args.n]
    if not imgs:
        sys.exit("No images found to evaluate.")

    print(f"Evaluating {len(imgs)} images · style={args.style} · stack={ACTIVE_STACK}", flush=True)
    options = get_vision_options("bn")

    # Precompute the (identical) prompt + encoded bytes per image.
    items = []
    for img in imgs:
        prompt = build_story_prompt(
            instruction=args.instruction,
            language="bn",
            style=args.style,
            child_name="",
            stack_key=ACTIVE_STACK,
            num_images=1,
        )
        items.append((img, prompt, [_encode(img)]))

    # Two phases so each serverless model cold-starts ONCE, not per image (the
    # alternating native→FT loop kept scaling the other model back to zero).
    print("Phase 1/2: native (base) ...", flush=True)
    natives = []
    for i, (img, prompt, b64) in enumerate(items, 1):
        t0 = time.time()
        story = (generate_story_remote(b64, prompt, options) or "").strip()
        dt = round(time.time() - t0, 1)
        print(f"  native [{i}/{len(items)}] {img.name}  {dt}s", flush=True)
        natives.append((story, dt))

    print("Phase 2/2: fine-tuned (vLLM) ...", flush=True)
    fts = []
    for i, (img, prompt, b64) in enumerate(items, 1):
        t0 = time.time()
        story = (generate_story_ft_remote(b64, prompt) or "").strip()
        dt = round(time.time() - t0, 1)
        print(f"  ft [{i}/{len(items)}] {img.name}  {dt}s", flush=True)
        fts.append((story, dt))

    rows = [
        (items[i][0], natives[i][0], natives[i][1], fts[i][0], fts[i][1])
        for i in range(len(items))
    ]

    out_dir = Path("finetune/eval_results")
    out_dir.mkdir(exist_ok=True)
    fname = out_dir / f"ft_vs_native_{datetime.now():%Y%m%d_%H%M}.md"
    lines = [
        f"# FT vs Native — Bengali story quality ({args.style})",
        f"Generated: {datetime.now():%Y-%m-%d %H:%M} · stack {ACTIVE_STACK} · {len(rows)} held-out images\n",
        "**Native** = base openbmb/MiniCPM-V-4_5 (current shipping Lever #1). ",
        "**FT** = merged Bengali LoRA via vLLM. Same prompt + image for both.\n",
        "> For the Bengali reviewer: for each image, which story reads more like a real "
        "grandmother's bedtime tale (natural words, রূপকথা imagery, no English/garbled "
        "words, calm sleepy ending)? Mark a winner per row.\n",
        "---\n",
    ]
    for img, native, tn, ft, tf in rows:
        lines += [
            f"## {img.name}",
            f"![{img.name}]({os.path.relpath(img, out_dir)})\n",
            f"### Native (base) — {tn}s",
            native or "_(empty)_", "",
            f"### FT (LoRA) — {tf}s",
            ft or "_(empty)_", "",
            "**Winner (reviewer):** ☐ Native  ☐ FT  ☐ Tie  ·  notes: ____",
            "\n---\n",
        ]
    fname.write_text("\n".join(lines))
    print(f"\nReport written to {fname}")
    print("Open it, have a Bengali speaker mark winners, and only set "
          "FINETUNED_VISION_MODEL if FT clearly wins.")


if __name__ == "__main__":
    main()