"""依頼110 F-2: 上カメラの日ごとの最初のコマで, 本線の四隅の検出 (mat_corner_detection.detect_corners) をかけ,
マットの上の端が画像の中に写っている日と, 画像の上の端で切れている日を分ける (読むだけ).

対象 (日ごとに1枚・それぞれの置き場の中で最も早いもの):
  runs/run_pick_place/  … recording/top.mp4 の最初のコマ, なければ上カメラの jpg (stage4_pick_tip → pick → place_pre → place)
  runs/pi05_policy/     … recording/top.mp4 の最初のコマ, なければ frames/0000.jpg (384×216)
  datasets/*/raw_episodes/ … frames/000000.jpg (384×216)
384×216 は 5 倍に広げてから検出する (本線の検出は 1920×1080 向けのしきい値).
切れている = 左上・右上の検出点の v が 6px 以下 (1920×1080 に直した値).
真っ暗なコマ (平均の明るさ 25 未満) は飛ばし, その日の次の試行を使う.

さらに, 1920×1080 のコマでは, マットに刷ってある 1cm の方眼の横線を左の端の内側で数え,
画像の上の端 (v=0) と, 上の2隅として検出された点が, マットの何 cm の線に当たるかを出す.
  方眼の「0」の線 = マットの下の端から約 10mm 上 (下の端は z=+361mm) → 「k cm」の線は z = 351 − 10k mm
  (450mm のマットに 0〜43cm の方眼. 本線の決めの上の端 z=−89mm は「43cm の線より 10mm 上」)
出力: f_days.csv / f_first_frames/{source}_{day}.jpg (検出点を描いたもの)
"""
import csv
import re
import sys
from pathlib import Path

import cv2
import numpy as np

sys.path.insert(0, "/Users/shimazueiichi/r2r2r-scripts")
from mat_corner_detection import detect_corners  # noqa: E402  本線の検出を読むだけで使う

ROOT = Path("/Volumes/KIOXIA/r2r2r-data")
OUT = ROOT / "lab" / "twin" / "20261107_失敗の再生と置き場所" / "F_上カメラの上の端"
IMG_DIR = OUT / "f_first_frames"
CUT_V = 6.0
DAY = re.compile(r"_(20\d{6})_(\d{6})")


def first_mp4_frame(p):
    cap = cv2.VideoCapture(str(p)); ok, f = cap.read(); cap.release()
    return f if ok else None


def candidates():
    """(source, day, stamp, path, kind) を返す."""
    out = []
    rp = ROOT / "runs" / "run_pick_place"
    for d in sorted(rp.iterdir()):
        m = DAY.search(d.name)
        if not m or not d.is_dir():
            continue
        mp4 = d / "recording" / "top.mp4"
        if mp4.exists():
            out.append(("run_pick_place", m.group(1), m.group(2), mp4, "mp4")); continue
        for lab in ("stage4_pick_tip", "pick", "place_pre", "place"):
            j = d / f"{d.name}_{lab}.jpg"
            if j.exists():
                out.append(("run_pick_place", m.group(1), m.group(2), j, f"jpg:{lab}")); break
    pp = ROOT / "runs" / "pi05_policy"
    for d in sorted(pp.iterdir()):
        m = DAY.search(d.name)
        if not m or not d.is_dir():
            continue
        mp4, j = d / "recording" / "top.mp4", d / "frames" / "0000.jpg"
        if mp4.exists():
            out.append(("pi05_policy", m.group(1), m.group(2), mp4, "mp4"))
        elif j.exists():
            out.append(("pi05_policy", m.group(1), m.group(2), j, "small"))
    for ds in sorted((ROOT / "datasets").iterdir()):
        re_ = ds / "raw_episodes"
        if not re_.is_dir():
            continue
        for e in sorted(re_.iterdir()):
            m = DAY.search(e.name); j = e / "frames" / "000000.jpg"
            if m and j.exists():
                out.append(("datasets", m.group(1), m.group(2), j, f"small:{ds.name}"))
    return out


def load(path, kind):
    if kind == "mp4":
        return first_mp4_frame(path), 1.0
    im = cv2.imread(str(path))
    if im is None:
        return None, 1.0
    if im.shape[1] < 1000:
        s = 1920 / im.shape[1]
        return cv2.resize(im, (1920, round(im.shape[0] * s)), interpolation=cv2.INTER_LINEAR), s
    return im, 1.0


def grid_lines(img, c):
    """左の端の内側の帯で方眼の横線を拾う. 戻り値: [(v, k cm)] (下の「0」の線から数えた番号) と周期 px."""
    (ulx, uly), (llx, lly) = c[0], c[3]
    vb = int(lly)
    prof = []
    for v in range(0, vb):
        xe = ulx + (llx - ulx) * (v - uly) / max(1e-6, (lly - uly))
        x0 = int(xe + 70); x1 = int(xe + 220)
        prof.append(img[v, x0:x1, 2].astype(float).mean() - img[v, x0:x1, 0].astype(float).mean() * 0.5)
    prof = np.array(prof)
    k = 31
    base = np.convolve(prof, np.ones(k) / k, mode="same")
    d = prof - base
    peaks = [v for v in range(2, len(d) - 2) if d[v] == d[v - 2:v + 3].max() and d[v] > 1.5]
    # 近すぎる山はまとめる
    merged = []
    for v in peaks:
        if merged and v - merged[-1] < 12:
            if d[v] > d[merged[-1]]:
                merged[-1] = v
            continue
        merged.append(v)
    if len(merged) < 10:
        return None, None
    sp = np.diff(merged)
    p = float(np.median(sp[(sp > 15) & (sp < 35)])) if np.any((sp > 15) & (sp < 35)) else None
    if p is None:
        return None, None
    # 下の端から 12px 以上上の, いちばん下の線を「0」とし, 上へ番号を振る
    lines = [v for v in merged if v < vb - 12]
    lines = lines[::-1]
    out = [(lines[0], 0)]
    for v in lines[1:]:
        kk = out[-1][1] + max(1, int(round((out[-1][0] - v) / p)))
        out.append((v, kk))
    return out, p


def cm_at(lines, v):
    """v の位置が何 cm の線に当たるか (線の間は直線で補う. 端の外は端の2本で延ばす)."""
    vs = np.array([a for a, _ in lines], float); ks = np.array([b for _, b in lines], float)
    o = np.argsort(vs); vs, ks = vs[o], ks[o]
    if v <= vs[0]:
        return float(ks[0] + (vs[0] - v) * (ks[0] - ks[1]) / (vs[1] - vs[0]))
    return float(np.interp(v, vs, ks))


def main():
    IMG_DIR.mkdir(exist_ok=True)
    seen, rows, skipped = set(), [], []
    for src, day, stamp, path, kind in candidates():
        if (src, day) in seen:
            continue
        img, s = load(path, kind)
        if img is None:
            continue
        if float(img.mean()) < 25:          # 真っ暗なコマ (カメラの映像が出ていない) はその日の次の試行を使う
            skipped.append(f"{src} {day} {path.name if hasattr(path, 'name') else path}: 暗い (平均 {img.mean():.0f})")
            continue
        c = detect_corners(img)
        row = {"source": src, "day": day, "first_run_time": stamp, "path": str(path).replace(str(ROOT) + "/", ""),
               "kind": kind, "scale_to_1920": round(s, 2)}
        if c is None:
            row["status"] = "四隅を検出できない"
            rows.append(row); seen.add((src, day)); continue
        c = np.asarray(c, float)
        row.update({f"{n}_{a}": round(float(c[i][j]), 1) for i, n in enumerate(("UL", "UR", "LR", "LL"))
                    for j, a in enumerate(("u", "v"))})
        top_v = min(c[0][1], c[1][1])
        row["top_cut"] = bool(top_v <= CUT_V)
        others = []
        if min(c[0][0], c[3][0]) <= CUT_V: others.append("左")
        if max(c[1][0], c[2][0]) >= img.shape[1] - 1 - CUT_V: others.append("右")
        if max(c[2][1], c[3][1]) >= img.shape[0] - 1 - CUT_V: others.append("下")
        row["other_cut"] = "・".join(others)
        if kind == "mp4" or kind.startswith("jpg"):
            lines, p = grid_lines(img, c)
            if lines:
                k0 = cm_at(lines, 0.0)
                kul = cm_at(lines, float(c[0][1]))
                row.update({"grid_period_px": round(p, 1), "grid_lines": len(lines),
                            "grid_top_line_cm": lines[-1][1],
                            "bottom_edge_to_0line_mm": round((c[3][1] - lines[0][0]) / p * 10, 1),
                            "image_top_cm": round(k0, 2), "image_top_z_mm": round(351 - 10 * k0, 1),
                            "UL_point_cm": round(kul, 2), "UL_point_z_mm": round(351 - 10 * kul, 1)})
        row["status"] = "切れている" if row["top_cut"] else "写っている"
        vis = img.copy()
        for i, (u, v) in enumerate(c):
            cv2.circle(vis, (int(u), int(v)), 14, (0, 255, 255), -1)
        cv2.putText(vis, f"{src} {day} {row['status'] == '切れている' and 'TOP CUT' or 'top visible'}",
                    (20, 60), cv2.FONT_HERSHEY_SIMPLEX, 1.6, (0, 255, 255), 4)
        cv2.imwrite(str(IMG_DIR / f"{src}_{day}.jpg"), cv2.resize(vis, (960, round(vis.shape[0] / 2))),
                    [cv2.IMWRITE_JPEG_QUALITY, 85])
        rows.append(row); seen.add((src, day))
    keys = []
    for r in rows:
        keys += [k for k in r if k not in keys]
    with open(OUT / "f_days.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); w.writerows(rows)
    print("飛ばしたコマ:", skipped)
    for r in rows:
        print(r["source"], r["day"], r["kind"][:20], r["status"], r.get("UL_v"), r.get("UR_v"), r.get("other_cut"),
              "top_z", r.get("image_top_z_mm"), "toplines", r.get("grid_top_line_cm"), "p", r.get("grid_period_px"),
              "bm", r.get("bottom_edge_to_0line_mm"))


if __name__ == "__main__":
    main()
