"""Diagnosis B (2026-10-13): can the base Qwen3.5-9B (no LoRA) recognize the target in our 320x180 frames?

Frames: GRPO v2 rollouts, rounds 1-5 (round 6 was cut off mid-rollout when training was paused), frames/NNNN.png =
the exact PNG bytes the policy saw. frames/i.png shows the pose after step i (steps.jsonl `pos`); frame 0 the
teleport location (train/grpo_locations.json) with yaw 0, pitch 0, which is also the origin of the spawn-relative
milestone targets. Angles follow benchmark_gen/utils.py (target yaw = atan2(-dx, dz)); a camera yaw delta adds to
pos.yaw (checked on round 1: 6346 of 6346 turns), so a target at a larger yaw is to the right. Pitch > 0 looks down.
- positives: target 3-12 blocks away, <= 20 deg left/right of the view center, <= 30 deg above/below it (the
  vertical half field of view is 35: Minecraft fov 70, options.txt fov:0.0);
- behind negatives: same distance band, >= 120 deg off center;
- wrong-name negatives: the positive frames, asked about a target name of another scene that is neither a target of
  this scene nor a block its build commands place.
At most 2 frames per episode and target, >= 20 frames apart, then round-robin over scenes.
Question "Is there a <name> in this image? Answer yes or no.", thinking off, temperature 0.
Writes items.jsonl, summary.md and contact_sheet.png (48 positives, red circle = projected target point).
Usage: diag_perception.py --url http://localhost:9041/v1 --model NAME [--n 200] [--out DIR]
"""
import argparse
import base64
from concurrent.futures import ThreadPoolExecutor
import json
import math
from pathlib import Path
import random

from openai import OpenAI
from PIL import Image, ImageDraw

ROOT = Path(__file__).resolve().parents[1]
GRPO = ROOT / "runs/grpo_v2_20261013"
ROUNDS = [f"round{r:03d}" for r in range(1, 6)]
BENCH = ROOT / "MineExplorer/benchmark"
LOCATIONS = json.loads((ROOT / "train/grpo_locations.json").read_text())
EYE = 1.62  # eye height above player_pos (the feet)
W, H, HALF_VFOV = 320, 180, 35.0


def episodes():
    """Finished rollout episodes: (episode dir, scene, steps, poses, targets); poses[i] is the pose of frames/i.png,
    targets are (milestone id, task, absolute target)."""
    for result in sorted(p for r in ROUNDS for p in (GRPO / r).glob("*/result.json")):
        ep = result.parent
        scene = ep.name.split("_")[0]
        x, y, z = LOCATIONS[scene]
        steps = [json.loads(line) for line in (ep / "steps.jsonl").read_text().splitlines()]
        poses = [{"x": x, "y": y, "z": z, "yaw": 0.0, "pitch": 0.0}] + [s["pos"] for s in steps]
        meta = json.loads((BENCH / scene / "multi-agent/metadata.json").read_text())
        targets = [(m["milestone_id"], m["task"], [t + o for t, o in zip(m["rules"][0]["params"]["target"], (x, y, z))])
                   for m in meta["milestones"]["milestones"]]
        yield ep, scene, steps, poses, targets


def view(pose, target):
    """(distance, horizontal, vertical): distance from the feet as the milestone check; horizontal = target yaw minus
    view yaw in [-180, 180), > 0 = right; vertical = target elevation from the eye above the view center."""
    dx, dy, dz = target[0] - pose["x"], target[1] - pose["y"] - EYE, target[2] - pose["z"]
    h = (math.degrees(math.atan2(-dx, dz)) - pose["yaw"] + 180) % 360 - 180
    v = math.degrees(math.atan2(dy, math.hypot(dx, dz))) + pose["pitch"]
    return math.dist((pose["x"], pose["y"], pose["z"]), target), h, v


def spread(candidates, n, rng):
    """At most 2 per (episode, target), >= 20 frames apart, then round-robin over scenes up to n."""
    rng.shuffle(candidates)
    groups = {}
    for c in candidates:
        g = groups.setdefault((c["ep"], c["mid"]), [])
        if len(g) < 2 and all(abs(c["frame"] - o["frame"]) >= 20 for o in g):
            g.append(c)
    by_scene = {}
    for c in (c for g in groups.values() for c in g):
        by_scene.setdefault(c["scene"], []).append(c)
    queues = [by_scene[s] for s in sorted(by_scene)]
    out = []
    while len(out) < n and any(queues):
        out += [q.pop() for q in queues if q][:n - len(out)]
    return out


def select(n):
    """(positives, behind negatives, wrong-name negatives), n of each where available."""
    rng = random.Random(0)
    pos, behind, names, placed = [], [], {}, {}
    for ep, scene, _, poses, targets in episodes():
        placed[scene] = " ".join(json.loads((BENCH / scene / "multi-agent/metadata.json").read_text())["commands"])
        for mid, task, target in targets:
            if not task.startswith("find "):  # a handful of "interact with X" / "open door" goals
                continue
            name = task[5:].removeprefix("a ").replace("_", " ")
            names.setdefault(scene, set()).add(name)
            for i, pose in enumerate(poses):
                dist, h, v = view(pose, target)
                item = {"ep": str(ep.relative_to(ROOT)), "scene": scene, "mid": mid, "frame": i, "name": name,
                        "asked": name, "dist": round(dist, 1), "h": round(h, 1), "v": round(v, 1)}
                if 3 <= dist <= 12 and abs(h) <= 20 and abs(v) <= 30:
                    pos.append(item)
                elif 3 <= dist <= 12 and abs(h) >= 120:
                    behind.append(item)
    pos, behind = spread(pos, n, rng), spread(behind, n, rng)
    every = sorted(set().union(*names.values()))
    wrong = [{**p, "asked": rng.choice([x for x in every if x not in names[p["scene"]]
                                        and x.replace(" ", "_") not in placed[p["scene"]]])} for p in pos]
    return pos, behind, wrong


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--n", type=int, default=200)
    ap.add_argument("--out", type=Path, default=ROOT / "runs/diag_20261013/perception")
    args = ap.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    pos, behind, wrong = select(args.n)
    print(f"positives {len(pos)} from {len({p['scene'] for p in pos})} scenes, behind {len(behind)} from "
          f"{len({p['scene'] for p in behind})} scenes, wrong-name {len(wrong)}", flush=True)

    client = OpenAI(api_key="EMPTY", base_url=args.url)

    def ask(kind_item):
        kind, item = kind_item
        png = (ROOT / item["ep"] / "frames" / f"{item['frame']:04d}.png").read_bytes()
        r = client.chat.completions.create(
            model=args.model, temperature=0, max_tokens=5,
            messages=[{"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": "data:image/png;base64," + base64.b64encode(png).decode()}},
                {"type": "text", "text": f"Is there a {item['asked']} in this image? Answer yes or no."}]}],
            extra_body={"chat_template_kwargs": {"enable_thinking": False}})
        answer = r.choices[0].message.content
        return {"kind": kind, **item, "answer": answer, "yes": answer.strip().lower().startswith("yes")}

    with ThreadPoolExecutor(16) as pool:
        items = list(pool.map(ask, [("pos", p) for p in pos] + [("behind", b) for b in behind]
                              + [("wrong", w) for w in wrong]))
    with (args.out / "items.jsonl").open("w") as f:
        f.writelines(json.dumps(item, ensure_ascii=False) + "\n" for item in items)

    rate = lambda xs: f"{100 * sum(x['yes'] for x in xs) / len(xs):.1f}%" if xs else "-"
    by_kind = {k: [x for x in items if x["kind"] == k] for k in ("pos", "behind", "wrong")}
    lines = ["| 类型 | 数量 | 答 yes 的比例 |", "|---|---:|---:|",
             f"| 正例（目标在视野中央附近） | {len(by_kind['pos'])} | {rate(by_kind['pos'])} |",
             f"| 背后负例（目标在身后） | {len(by_kind['behind'])} | {rate(by_kind['behind'])} |",
             f"| 错名负例（同一批正例帧，问别的场景的目标名） | {len(by_kind['wrong'])} | {rate(by_kind['wrong'])} |"]
    per = sorted({x["asked"] for x in items}, key=lambda n: -sum(x["asked"] == n for x in by_kind["pos"]))[:20]
    lines += ["", "逐目标（正例最多的 20 个；错名一列是把这个名字问在不含它的帧上）：", "",
              "| 目标 | 正例 n | 正例 yes | 背后 n | 背后 yes | 错名 n | 错名 yes |", "|---|---:|---:|---:|---:|---:|---:|"]
    for n in per:
        cells = [[x for x in by_kind[k] if x["asked"] == n] for k in ("pos", "behind", "wrong")]
        lines.append(f"| {n} | " + " | ".join(f"{len(c)} | {rate(c)}" for c in cells) + " |")
    (args.out / "summary.md").write_text(
        "帧来自 GRPO v2 第 1 到 5 轮 rollout（策略实际看到的 320×180 PNG）。正例：目标距离 3 到 12 格，偏离视野中心左右 ≤20°、"
        "上下 ≤30°（竖直半视场 35°）；背后负例：同距离段，偏离 ≥120°；错名负例：同一批正例帧，问另一个场景的目标名"
        "（本场景的目标和建造命令里都没有这种方块，但个别名字如南瓜、甘蔗可能在自然地形里出现）。每个回合每个目标最多 2 帧、"
        f"相隔 ≥20 帧，按场景轮流抽：正例来自 {len({p['scene'] for p in pos})} 个场景，背后负例来自 "
        f"{len({p['scene'] for p in behind})} 个场景。只按几何选帧，不判断遮挡；"
        "接触表 contact_sheet.png 是前 48 个正例（红圈为目标点按视角投影的位置，标注为目标名、回答、距离、回合/帧号），"
        "全部逐条结果见 items.jsonl。\n\n" + "\n".join(lines) + "\n")

    sheet = Image.new("RGB", (6 * W, 8 * (H + 14)), "white")
    for k, x in enumerate(by_kind["pos"][:48]):
        img = Image.open(ROOT / x["ep"] / "frames" / f"{x['frame']:04d}.png").convert("RGB")
        draw = ImageDraw.Draw(img)
        # approximate pinhole projection from the two angles (exact at pitch 0); 16:9, vertical fov 70
        u = W / 2 + W / 2 * math.tan(math.radians(x["h"])) / (math.tan(math.radians(HALF_VFOV)) * W / H)
        v = H / 2 - H / 2 * math.tan(math.radians(x["v"])) / math.tan(math.radians(HALF_VFOV))
        draw.ellipse((u - 6, v - 6, u + 6, v + 6), outline="red", width=2)
        cell = (k % 6 * W, k // 6 * (H + 14))
        sheet.paste(img, (cell[0], cell[1] + 14))
        ImageDraw.Draw(sheet).text((cell[0] + 2, cell[1] + 1), f"{x['name']} | {x['answer'].strip()[:5]} | {x['dist']:.0f}m | "
                                   f"{Path(x['ep']).name}/{x['frame']}", fill="black")
    sheet.save(args.out / "contact_sheet.png")
    print((args.out / "summary.md").read_text(), flush=True)


if __name__ == "__main__":
    main()