#!/usr/bin/env python3
"""过夜编排：等 arope_5k_ga2（new_weight 2）训完 → 评测 → 训 new_weight 5 的同配置 5k 步 → 评测 → 写对比汇总。

只做三件事：轮询日志、subprocess 跑本仓库自己的脚本、写 outputs/overnight/ 下的汇总。任一阶段失败不中断后续阶段，
错误写进 outputs/overnight/status.json。

    nohup .venv/bin/python scripts/overnight.py > outputs/overnight/overnight.log 2>&1 &
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from datetime import datetime

ROOT = "/opt/dlami/nvme/zhiyangdeng/ActionRoPE"
PY = f"{ROOT}/.venv/bin/python"
OUT = f"{ROOT}/outputs/overnight"
ENV = dict(os.environ, DIFFSYNTH_SKIP_DOWNLOAD="True", PYTHONPATH=ROOT)
RUNS = [
    # (run 名, 说明, 额外训练参数, 是否需要本脚本来训)
    ("arope_5k_ga2", "flat ×2", [], False),      # 已在跑，只等它结束
    # 第 5 点「权重形状微调」：new 区权重按离首帧足迹边界的距离从 1 爬到 3（4 latent 格），再按噪声水平 ×(0.5+σ)
    ("arope_5k_ga2_shape", "dist 1→3 ramp4 + σ", ["--new_weight", "3.0", "--new_weight_shape", "dist",
                                                  "--new_weight_ramp", "4.0", "--new_weight_sigma"], True),
]
STATUS = {"started": datetime.now().isoformat(timespec="seconds"), "stages": {}}


def log(msg):
    print(f"[{datetime.now():%H:%M:%S}] {msg}", flush=True)


def save_status():
    os.makedirs(OUT, exist_ok=True)
    with open(f"{OUT}/status.json", "w", encoding="utf-8") as f:
        json.dump(STATUS, f, ensure_ascii=False, indent=1)


def stage(name, fn):
    if STATUS["stages"].get(name, {}).get("ok"):
        log(f"== {name} 已完成，跳过")
        return
    t0 = time.time()
    log(f"== {name} 开始")
    try:
        r = fn()
        STATUS["stages"][name] = {"ok": True, "sec": round(time.time() - t0), "result": r}
        log(f"== {name} 完成 {time.time() - t0:.0f}s")
    except Exception as e:  # 不让一个阶段拖死整晚
        STATUS["stages"][name] = {"ok": False, "sec": round(time.time() - t0), "error": repr(e)}
        log(f"!! {name} 失败: {e!r}")
    save_status()


def run(cmd, log_path, env_extra=None, timeout=None):
    env = dict(ENV, **(env_extra or {}))
    with open(log_path, "a", encoding="utf-8") as lf:
        lf.write(f"\n$ {' '.join(cmd)}\n")
        lf.flush()
        rc = subprocess.call(cmd, cwd=ROOT, env=env, stdout=lf, stderr=subprocess.STDOUT, timeout=timeout)
    if rc != 0:
        raise RuntimeError(f"rc={rc}: {' '.join(cmd)}  (log {log_path})")


def wait_train_done(run_name, pid_file=None, max_hours=6.0):
    """轮询 train.log 出现 [done] 且训练进程退出。"""
    log_path = f"{ROOT}/outputs/{run_name}/train.log"
    t0 = time.time()
    while True:
        done = os.path.exists(log_path) and "[done]" in open(log_path, encoding="utf-8", errors="ignore").read()[-4000:]
        alive = False
        if pid_file and os.path.exists(pid_file):
            pid = int(open(pid_file).read().strip() or 0)
            alive = False
            if pid > 0 and os.path.exists(f"/proc/{pid}"):
                # 僵尸进程（父进程还没 wait）在 /proc 里仍然存在，不能算活着，否则会永远等下去（2026-09-05 踩过）
                with open(f"/proc/{pid}/stat") as fh:
                    alive = fh.read().split(")")[-1].split()[0] != "Z"
        if done and not alive:
            return {"train_log": log_path, "waited_sec": round(time.time() - t0)}
        if os.path.exists(log_path):
            tail = open(log_path, encoding="utf-8", errors="ignore").read()[-3000:]
            if "Traceback" in tail and not alive:
                raise RuntimeError(f"{run_name} 训练异常退出，见 {log_path}")
        if time.time() - t0 > max_hours * 3600:
            raise RuntimeError(f"{run_name} 等待超过 {max_hours} h")
        time.sleep(30)


def train(run_name, extra_args):
    """用 scripts/train_arope.sh（GA=2）起一个 5k 步训练并等它结束。"""
    if os.path.exists(f"{ROOT}/outputs/{run_name}/step-5000.safetensors"):
        return {"skipped": "step-5000.safetensors 已存在"}
    launch_log = f"{ROOT}/outputs/{run_name}_launch.log"
    pid_file = f"{ROOT}/outputs/{run_name}.pid"
    env = dict(ENV, RUN=run_name, GA="2")
    with open(launch_log, "a", encoding="utf-8") as lf:
        p = subprocess.Popen(["bash", f"{ROOT}/scripts/train_arope.sh", "--max_steps", "5000", "--warmup_steps", "200",
                              *extra_args], cwd=ROOT, env=env, stdout=lf, stderr=subprocess.STDOUT)
    open(pid_file, "w").write(str(p.pid))
    log(f"{run_name} 已启动 pid={p.pid} 参数={extra_args}")
    r = wait_train_done(run_name, pid_file)
    p.wait(timeout=600)
    r["rc"] = p.returncode
    return r


def evaluate(run_name):
    ck = f"{ROOT}/outputs/{run_name}/step-5000.safetensors"
    if not os.path.exists(ck):
        raise FileNotFoundError(ck)
    res = {}
    # 四科（bg50，全部 1×）
    run([PY, "-m", "actionrope.eval_bg50", "--ckpt", ck, "--run", run_name, "--stages", "gen,measure"],
        f"{OUT}/eval_{run_name}.log", timeout=3 * 3600)
    s = json.load(open(f"{ROOT}/outputs/eval_bg50/{run_name}/summary.json", encoding="utf-8"))
    res["bg50_summary"] = s
    # 训练场景首帧的 8 条演示（含 up/down，用来看纵向增益 1.12 之后还超不超调）
    run(["bash", f"{ROOT}/scripts/eval_5k_demo.sh"], f"{OUT}/demo_{run_name}.log",
        env_extra={"CK": ck, "OUT": f"{ROOT}/outputs/samples/{run_name}"}, timeout=3600)
    demo = {}
    d = f"{ROOT}/outputs/samples/{run_name}"
    for name in ("replay", "right_x1", "left_x1", "up_x1", "down_x1", "right_x0.5", "right_x1.5", "there_back"):
        p = f"{d}/{name}.json"
        if os.path.exists(p):
            j = json.load(open(p, encoding="utf-8"))
            demo[name] = {"cmd_bg_shift_80": [-v for v in j["frame_offset_px_80"]],
                          "measured": (j.get("measured") or {}).get("sift"),
                          "psnr_mean_1_80": j.get("psnr_mean_1_80")}
    res["demo"] = demo
    return res


def pick(summary, *keys, default=None):
    cur = summary
    for k in keys:
        if not isinstance(cur, dict) or k not in cur:
            return default
        cur = cur[k]
    return cur


def write_summary():
    rows = []
    for run_name, desc, _extra, _need in RUNS:
        st = STATUS["stages"].get(f"eval:{run_name}", {})
        r = st.get("result") or {}
        rows.append((run_name, desc, r))
    lines = ["# 过夜对比：new 区权重 flat ×2 vs 形状化（dist 1→3 + σ 调制）；同配置：8 卡 × GA2、5k 步、lr 1e-5、纵向增益 1.12", "",
             f"生成于 {datetime.now():%Y-%m-%d %H:%M}。旧的 arope_5k（GA1、纵向增益 1.0、new_weight 2）四科在 README 里。", "",
             "## bg50 四科（全部 1×）", "",
             "| 指标 | " + " | ".join(f"{n} ({w})" for n, w, _ in rows) + " |",
             "|---|" + "---|" * len(rows)]
    metrics = [
        ("增益线性 斜率", ("gain", "slope")), ("增益线性 R²", ("gain", "r2")),
        ("折返 PSNR 中位", ("back", "psnr_med")), ("折返 SSIM 中位", ("back", "ssim_med")), ("折返 NCC 中位", ("back", "ncc_med")),
        ("折返 回归残差中位 px", ("back", "resid_px_med")), ("折返 补偿后 NCC", ("back", "ncc_comp_med")),
        ("倒退帧总数", ("smooth", "backward_frames_total")), ("速度波动 std/mean 中位", ("smooth", "speed_cv_med")),
        ("平均速度/承诺", ("smooth", "speed_ratio_med")),
        ("角度误差中位 °", ("direction", "angle_abs_med")), ("角度误差 p90 °", ("direction", "angle_abs_p90")),
        ("幅度比中位", ("direction", "ratio_med")), ("失控数", ("direction", "n_fail")),
    ]
    for label, keys in metrics:
        vals = []
        for _, _, r in rows:
            v = pick(r.get("bg50_summary", {}), *keys)
            if v is None:  # summary 的键名以实际文件为准；找不到就把整个 summary 留给人看
                v = "?"
            vals.append(f"{v:.3f}" if isinstance(v, float) else str(v))
        lines.append(f"| {label} | " + " | ".join(vals) + " |")
    lines += ["", "（键名对不上的格子显示 ?，完整数字见 outputs/eval_bg50/<run>/summary.json 与 report.md）", "",
              "## 训练场景首帧演示：指令 vs 实测背景位移（帧 0→80，px）", ""]
    for run_name, w, r in rows:
        lines.append(f"### {run_name} ({w})")
        lines.append("| 指令 | 指令位移 | 实测 | 帧1–80 PSNR |")
        lines.append("|---|---|---|---|")
        for name, d in (r.get("demo") or {}).items():
            c = d["cmd_bg_shift_80"]; m = d["measured"]
            lines.append(f"| {name} | ({c[0]:+.0f},{c[1]:+.0f}) | " + (f"({m[0]:+.0f},{m[1]:+.0f})" if m else "测不出")
                         + " | " + (f"{d['psnr_mean_1_80']:.1f}" if d.get("psnr_mean_1_80") else "—") + " |")
        lines.append("")
    lines += ["## 阶段状态", "", "```", json.dumps(STATUS["stages"], ensure_ascii=False, indent=1, default=str)[:6000], "```"]
    with open(f"{OUT}/summary.md", "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    return f"{OUT}/summary.md"


def main():
    os.makedirs(OUT, exist_ok=True)
    if os.path.exists(f"{OUT}/status.json") and "--resume" in sys.argv:
        STATUS.update(json.load(open(f"{OUT}/status.json", encoding="utf-8")))
    save_status()
    for run_name, _desc, extra, need_train in RUNS:
        if need_train:
            stage(f"train:{run_name}", lambda: train(run_name, extra))
        else:
            stage(f"wait:{run_name}", lambda: wait_train_done(run_name, f"{ROOT}/outputs/{run_name}.pid"))
        stage(f"eval:{run_name}", lambda: evaluate(run_name))
    stage("summary", write_summary)
    STATUS["finished"] = datetime.now().isoformat(timespec="seconds")
    save_status()
    log("全部结束")


if __name__ == "__main__":
    main()