#!/usr/bin/env python3
"""teaser · 世界编辑的底子挑选 + 编辑计划（CPU）→ <out>/bases.jsonl、<out>/plan.jsonl（edits.py run 的输入）。

底子 = 已有 rollout（outputs/teaser/roll/<name>/）里：首帧 A 档（prod/picks.jsonl）、vscore 无硬伤（flicker 不算）、变体 ∈ {原首帧, rc}、
teaser 分 ≥ --min_score；贪心按 teaser 分取 --n 个，同一张地图（map case）只取一条、每个主题 ≤ --per_theme、每种画风 ≤ --per_look、
同一个场景 ≤ 2 ⇒ 主题 / 场景 / 画风都分散。另算每条底子的「Boss 显眼度」（帧 24..80 Boss 像素中位）和「玩家走了多远」（帧 16..80），
决定给它排哪些编辑：
  * 人人都有：add ×2（自动选反差大的物体，k 32 / 40）、grow（k 24，dur 16）、delete（pop / sink 交替，k 32）、swap（wave 16，k 28）、
    write（马良 / MALIANG / ★ / ♥ 轮换，k 20，dur 32；每第 3 条用 raise 立体字）、spawn（树 / 柱 / 立石 / 水晶轮换，k 24）、
    delete class（k 36，每第 2 条）；
  * 玩家走得够远（≥ 3 m）：paint（k 16；每第 2 条 lead 16 = 路铺在前面）；
  * Boss 显眼（像素中位 ≥ 1500）：boss_del（k 32）、boss_clone（k 32）、trap（k 32）。

usage：.venv/bin/python scripts/teaser/edit_plan.py [--n 48] [--out outputs/teaser/edit/e81] [--exclude a,b]
       .venv/bin/python scripts/teaser/edit_plan.py --wave2 outputs/teaser/edit/e81/plan2.jsonl   # 第二波（只做立体物体类编辑）
       .venv/bin/python scripts/teaser/edit_plan.py --wave3 outputs/teaser/edit/e81/plan3.jsonl   # 第三波（vegetation 地面图层）
"""
from __future__ import annotations

import argparse
import collections
import json
import os

import numpy as np

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
T = f"{M}/outputs/teaser"
HARD = {"scene_lost", "degraded", "broken", "char_melt", "blur"}


def candidates(min_score: float) -> list[dict]:
    tier = {}
    for ln in open(f"{T}/prod/picks.jsonl"):
        r = json.loads(ln)
        tier[r["case"] + "__" + r["scene"]] = r.get("tier")
    rows = []
    for f in os.listdir(f"{T}/vscore/roll"):
        name = f[:-len(".vscore.json")]
        parts = name.split("__")
        var = parts[5] if len(parts) > 5 else ""
        if var not in ("", "rc"):
            continue
        base = "__".join(parts[:5])
        if tier.get(base) != "A":
            continue
        v = json.load(open(f"{T}/vscore/roll/{f}"))
        if set(v.get("flags") or []) & HARD or v["teaser"] < min_score:
            continue
        if not os.path.exists(f"{T}/roll/{name}/meta.json"):
            continue
        rows.append(dict(name=name, case=base.rsplit("__", 1)[0], theme=parts[0], scene=parts[4].split("@")[0],
                         look=parts[4].split("@")[1], var=var, teaser=v["teaser"], flags=v.get("flags") or []))
    return sorted(rows, key=lambda r: -r["teaser"])


def pick(rows, n, per_theme, per_look, exclude=()):
    out, seen_case = [], set()
    th, lk, sc = collections.Counter(), collections.Counter(), collections.Counter()
    for r in rows:
        if r["name"] in exclude or r["case"] in seen_case:
            continue
        if th[r["theme"]] >= per_theme or lk[r["look"]] >= per_look or sc[r["scene"]] >= 2:
            continue
        out.append(r)
        seen_case.add(r["case"])
        th[r["theme"]] += 1
        lk[r["look"]] += 1
        sc[r["scene"]] += 1
        if len(out) >= n:
            break
    return out


def stats(r):
    import sys
    sys.path.insert(0, f"{M}/code/gc25d")
    ids = np.load(f"{T}/roll/{r['name']}/ids.npz")["ids"]
    boss = ((ids >= 17) & (ids <= 21)) | (ids == 11)
    bpx = boss[24:].sum((1, 2))
    meta = json.load(open(f"{T}/roll/{r['name']}/meta.json"))
    src = json.load(open(os.path.join(os.path.dirname(meta["job"]["ids"]) if os.path.isabs(meta["job"]["ids"])
                                      else os.path.join(M, os.path.dirname(meta["job"]["ids"])), "source.json")))
    m = np.load(f"{M}/data/gc25d_v10/meta/{src['clip_id']}.npz")
    P = m["actor_pos"][16:81, 0, :2]
    walk = float(np.linalg.norm(np.diff(P, axis=0), axis=1).sum()) if np.isfinite(P).all() else 0.0
    return dict(boss_px_med=int(np.median(bpx)), walk_cm=round(walk, 1), ortho=bool(src["ortho"]), clip_id=src["clip_id"])


TEXTS = ["马良", "MALIANG", "star", "马良", "heart", "HELLO", "马良", "GG"]
SPAWN = ["tree", "pillar", "stone", "crystal"]


def plan_rows(bases):
    rows = []
    for i, b in enumerate(bases):
        nm = b["name"]
        add = lambda e, k, p=None, prio=0: rows.append(dict(base=nm, edit=e, k=k, params=p or {}, prio=prio))
        add("add", 32, {}, 2)
        add("add", 40, {"size": 1.15}, 1)
        add("grow", 24, {"dur": 16}, 2)
        add("delete", 32, {"mode": "pop" if i % 2 == 0 else "sink"}, 1)
        add("swap", 28, {"wave": 16}, 2)
        t = TEXTS[i % len(TEXTS)]
        add("write", 20, {"text": t, "dur": 32, "mode": "raise" if i % 3 == 2 else "paint"}, 2)
        add("spawn", 24, {"obj": SPAWN[i % len(SPAWN)]}, 1)
        if i % 2 == 1:
            add("delete", 36, {"mode": "class"}, 0)
        if b["walk_cm"] >= 300:
            add("paint", 16, {"lead": 16} if i % 2 else {}, 2)
        if b["boss_px_med"] >= 1500:
            add("boss_del", 32, {}, 1)
            add("boss_clone", 32, {}, 1)
            add("trap", 32, {}, 1)
    return rows


ADD2 = ["tower", "house", "statue", "crystal", "tree", "wall", "pillar", "stone"]
GROW2 = ["wall", "crystal", "tree", "stone", "tower"]
DROP = ["props", "boulder"]
TRAIL = ["bush", "crystal", "stone", "pillar"]
SWAP2 = ["veg2rock", "arch2veg", "rock2veg", "arch2rock", "veg2arch"]
RAISE = [("star", "arch"), ("HI", "veg"), ("heart", "veg"), ("GO", "arch"), ("arrow", "veg"), ("WIN", "arch")]


def plan_rows2(bases):
    """第二波（只做立体物体类的编辑：第一波看下来地面图层类 —— paint / write paint / 地面 swap —— 模型基本不画）。"""
    rows = []
    for i, b in enumerate(bases):
        nm = b["name"]
        add = lambda e, k, p=None, prio=0: rows.append(dict(base=nm, edit=e, k=k, params=p or {}, prio=prio))
        add("add", 24, {"obj": ADD2[i % len(ADD2)]}, 1)
        add("add", 48, {"obj": ADD2[(i + 3) % len(ADD2)]}, 0)
        add("add", 32 + 8 * (i % 2), {"anim": "drop", "obj": DROP[i % 2]}, 1)
        add("grow", 32, {"obj": GROW2[i % len(GROW2)], "dur": 20}, 1)
        add("delete", 32 if i % 2 == 0 else 40, {"mode": "sink" if i % 2 == 0 else "pop"}, 0)
        if i % 2 == 0:
            add("delete", 36, {"mode": "class"}, 1)
        for j in range(2):
            add("swap", 28, {"pair": SWAP2[(i + j) % len(SWAP2)], "wave": 16}, 1)
        add("swap", 36, {"pair": SWAP2[(i + 2) % len(SWAP2)]}, 0)
        add("spawn", 20, {"obj": TRAIL[i % len(TRAIL)] if TRAIL[i % len(TRAIL)] != "bush" else "tree", "mode": "avenue"}, 1)
        add("spawn", 24, {"obj": ["tree", "pillar", "stone", "crystal"][(i + 1) % 4], "loose": 1}, 0)
        add("paint", 16, {"mode": "trail3d", "obj": TRAIL[(i + 1) % len(TRAIL)]}, 1)
        t, c = RAISE[i % len(RAISE)]
        add("write", 20, {"text": t, "mode": "raise", "cls": c, "big": 1.3, "shrink": 0.7, "cell": 10, "max_run": 1,
                          "bold": 0.6 if t.isalpha() and t.isupper() else 0.0}, 1)
        if b["boss_px_med"] >= 1500:
            add("trap", 32, {"min_keep": 0.4, "cls": ["arch", "veg", "rock"][i % 3]}, 1)
    return rows


def plan_rows3(bases):
    """第三波：地面图层里只有 vegetation 画得出来（铺装上的草 / 花径）⇒ 画路 / 写字 / 地面 swap 都改用 vegetation。"""
    rows = []
    texts = ["马良", "star", "heart", "MALIANG", "马良", "HI"]
    for i, b in enumerate(bases):
        nm = b["name"]
        add = lambda e, k, p=None, prio=0: rows.append(dict(base=nm, edit=e, k=k, params=p or {}, prio=prio))
        if b["walk_cm"] >= 300:
            add("paint", 16, {"cls": "veg", "lead": 16}, 2)
            add("paint", 20, {"cls": "veg"}, 1)
        add("write", 20, {"text": texts[i % len(texts)], "cls": "veg", "big": 1.2, "dur": 32}, 2)
        add("swap", 28, {"pair": "pave2veg", "wave": 20}, 1)
        add("swap", 28, {"pair": "terrain2veg", "wave": 20}, 1)
    return rows


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=48)
    ap.add_argument("--min_score", type=float, default=0.8)
    ap.add_argument("--per_theme", type=int, default=3)
    ap.add_argument("--per_look", type=int, default=12)
    ap.add_argument("--out", default=f"{T}/edit/e81")
    ap.add_argument("--exclude", default="")
    ap.add_argument("--append", action="store_true", help="追加到已有 plan（新底子）")
    ap.add_argument("--wave2", default="", help="给 <out>/bases.jsonl 里的底子出第二波计划 → 这个路径")
    ap.add_argument("--wave3", default="", help="第三波（vegetation 地面图层）→ 这个路径")
    a = ap.parse_args(argv)
    if a.wave2 or a.wave3:
        bases = [json.loads(l) for l in open(f"{a.out}/bases.jsonl")]
        pr = plan_rows2(bases) if a.wave2 else plan_rows3(bases)
        with open(a.wave2 or a.wave3, "w") as fh:
            for r in pr:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"{len(bases)} 个底子 · {len(pr)} 行 → {a.wave2 or a.wave3}")
        print("编辑", dict(collections.Counter(r["edit"] for r in pr)))
        return
    rows = candidates(a.min_score)
    ex = set(x for x in a.exclude.split(",") if x)
    if a.append and os.path.exists(f"{a.out}/bases.jsonl"):
        ex |= {json.loads(l)["name"] for l in open(f"{a.out}/bases.jsonl")}
    bases = pick(rows, a.n, a.per_theme, a.per_look, ex)
    for b in bases:
        b.update(stats(b))
    os.makedirs(a.out, exist_ok=True)
    mode = "a" if a.append else "w"
    with open(f"{a.out}/bases.jsonl", mode) as fh:
        for b in bases:
            fh.write(json.dumps(b, ensure_ascii=False) + "\n")
    pr = plan_rows(bases)
    with open(f"{a.out}/plan.jsonl", mode) as fh:
        for r in pr:
            fh.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"{len(rows)} 个候选 → {len(bases)} 个底子 · {len(pr)} 行计划")
    print("主题", dict(collections.Counter(b["theme"] for b in bases)))
    print("画风", dict(collections.Counter(b["look"] for b in bases)))
    print("编辑", dict(collections.Counter(r["edit"] for r in pr)))


if __name__ == "__main__":
    main()