#!/usr/bin/env python3
"""camera_align.py — 上カメラの写り方をお手本に合わせる画面 (Mac 側・アームは動かさない).

上カメラの iPhone は普段使いのもので, 毎回付け外しする. 付けるたびに写り方が変わり,
2026-10-21 の実機では お手本より下に大きくずれたまま動かしてしまった
(runs/pi05_policy/pi05_20261021_082645 は dy=+37.5. 1 分前のドライランは dy=-7.8 で正常).
iPhone を付けたあと, この画面で写り方を合わせてから実機に進む.

使い方は 2 つ:
  1. 上カメラの写り方を お手本と同じに戻す (dy / dx / response を見る)
  2. 重なったお手本に写っているブロックに実物のブロックを合わせて, お手本と同じ位置に置く
     (--episode でどのお手本を重ねるか選ぶ. ブロックの位置はエピソードごとに違う)

画面:
  ライブ映像に 基準画像 (--episode のエピソードの 1 枚目・HOME の姿勢) を半透明で重ねる.
  --episode を省くと --raw-dir の最初のエピソードになる (従来どおり).
  隅に写り方のずれ (dx, dy) と相関の強さ (response) を出し, 判定を大きく表示する:
    OK       しきい値以内
    ZURE     dx か dy がしきい値の外 (= ずれ)
    NO MATCH 相関が弱く dx/dy が当てにならない (= 測れない)
  q または Esc で終了. **表示だけで, ファイルには何も書かない.**

run_policy_client.py は開始時に この measure_framing / framing_verdict で同じ検査をして,
OK 以外ならアームに何も送らずに終了する (--ignore-framing で続行).

前提:
  - 上カメラ (iPhone 連係) が見えている. アームは HOME に置いておく (基準画像が HOME の姿勢)

Usage:
  RAW=/Volumes/KIOXIA/r2r2r-data/datasets/mecharm_redblock_box_teach_v1_trim2/raw_episodes
  python3 pi05/camera_align.py --raw-dir $RAW
  python3 pi05/camera_align.py --raw-dir $RAW --episode run_0001_20261019_192343  # 位置を変える
  python3 pi05/camera_align.py --raw-dir $RAW --episode 192343                    # 部分一致でもよい
  python3 pi05/camera_align.py --raw-dir $RAW --alpha 0.7 --scale 2   # 基準を濃く / 画面を小さく
"""
import argparse
import sys
import time
from pathlib import Path

import cv2
import numpy as np

_HERE = Path(__file__).resolve().parent
_ROOT = _HERE.parent
_PLACE = _ROOT / "dynamic_tracking_limits" / "place" / "place_modules"   # camera_resolver はここ
for p in (_HERE, _ROOT, _PLACE):
    sys.path.insert(0, str(p))

import camera_resolver  # noqa: E402

IMG_LONG = 384            # build_demo_episodes.py / run_policy_client.py と同じ (学習画像の長辺)
FRAMING_MAX_DY = 20.0     # |dy| がこれを超えたら「ずれ」. 下の実測値の間を取った
FRAMING_MAX_DX = 20.0     # |dx| の上限 (codex R2). お手本の dx は +0.00..+7.23 (2026-10-21 実測)
# response (相関の強さ) がこれ未満なら, dx/dy を信用できないので「測れない」として止める (codex R1).
# 2026-10-21 実測: 正しく対応している比較は 0.411..1.000.
# 無関係な画像は ランダム 0.0375 / 真っ黒 0.0000, 10/17 の 222723 の 1 枚目は 0.0505
# (dy=+8.2 / dx=-2.5 なので, 相関を見ないと「OK」で通ってしまう).
FRAMING_MIN_RESPONSE = 0.20
WINDOW = "camera_align (q で終了)"


def to_small_bgr(frame):
    """学習画像と同じ大きさに縮める (run_policy_client.grab_rgb と同じ処理)."""
    h, w = frame.shape[:2]
    s = IMG_LONG / max(h, w)
    return cv2.resize(frame, (round(w * s), round(h * s)), interpolation=cv2.INTER_AREA)


def pick_episode(raw_dir: Path, episode=None):
    """基準にするエピソードのフォルダを決める. episode=None なら最初のエピソード (従来どおり).

    episode は エピソード名 (例 run_0001_20261019_192343). 完全一致が無ければ,
    1 本だけに当てはまる部分一致 (例 192343) も受け付ける.
    """
    eps = sorted(p.parent for p in raw_dir.glob("*/episode.npz"))
    if not eps:
        sys.exit(f"[ERROR] お手本が見つからない (基準画像を決められない): {raw_dir}")
    if episode is None:
        return eps[0]
    exact = [p for p in eps if p.name == episode]
    if exact:
        return exact[0]
    part = [p for p in eps if episode in p.name]
    if len(part) == 1:
        return part[0]
    names = "\n".join(f"           {p.name}" for p in eps)
    if len(part) > 1:
        sys.exit(f"[ERROR] --episode {episode} が {len(part)} 本に当てはまります:\n"
                 + "\n".join(f"           {p.name}" for p in part))
    sys.exit(f"[ERROR] --episode {episode} が見つからない. {raw_dir} にあるのは {len(eps)} 本:\n{names}")


def grasp_j0(ep_dir: Path):
    """つかんだ瞬間 (グリッパーが閉に変わったステップ) の j0 [deg]. 分からなければ None.

    どのブロック位置のお手本かを表示するために使う (check_dataset.py の項目 5 と同じ見方).
    """
    try:
        s = np.load(ep_dir / "episode.npz")["state"]
    except Exception:
        return None
    ch = np.flatnonzero(s[1:, 6] > s[:-1, 6]) + 1
    return float(s[ch[0], 0]) if len(ch) else None


def load_reference(raw_dir: Path, episode=None):
    """基準画像 = エピソードの 1 枚目 (HOME の姿勢). 戻り値: (BGR 画像, パス).

    episode=None (既定) なら --raw-dir の最初のエピソード. run_policy_client.py は
    こちらで呼ぶので, 写り方の検査に使う基準は常に最初のエピソードのまま.
    """
    ep = pick_episode(raw_dir, episode)
    frames = sorted((ep / "frames").glob("*.jpg"))
    if not frames:
        sys.exit(f"[ERROR] お手本の画像が無い: {ep}")
    bgr = cv2.imread(str(frames[0]))
    if bgr is None:
        sys.exit(f"[ERROR] 基準画像を読めない: {frames[0]}")
    return bgr, frames[0]


def measure_framing(ref_bgr, cur_bgr):
    """写り方のずれを測る. 戻り値: (dx, dy, response). cv2.phaseCorrelate + ハニング窓.

    2026-10-21 実測 (基準 = teach_v1 の run_0001_20261019_183454 の 1 枚目):
      お手本 191257   dy= +9.1  / お手本 195120 dy= -6.7   ← ブロックの位置で揺れる範囲
      ドライラン 082512 dy= -7.8                            ← お手本と同じ写り方. 通す
      実機 082645       dy=+37.5                            ← 下にずれた写り方. 止める
    dx はブロックの位置で ±6 ほど動く (お手本どうしで +5.0 / +5.8) ので判定には使わず, 表示だけ.
    """
    a = cv2.cvtColor(ref_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)
    b = cv2.cvtColor(cur_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)
    if a.shape != b.shape:
        b = cv2.resize(b, (a.shape[1], a.shape[0]), interpolation=cv2.INTER_AREA)
    win = cv2.createHanningWindow((a.shape[1], a.shape[0]), cv2.CV_32F)
    (dx, dy), response = cv2.phaseCorrelate(a, b, win)
    return float(dx), float(dy), float(response)


def framing_verdict(dx, dy, response):
    """写り方の判定. 戻り値: (verdict, 理由). verdict は "OK" / "ずれ" / "測れない".

    相関が弱い (response が低い) ときは dx/dy 自体が当てにならないので,
    「ずれ」ではなく「測れない」として扱う (既定ではどちらも止める).
    """
    if not all(np.isfinite(v) for v in (dx, dy, response)):
        return "測れない", "dx / dy / response に非有限値"
    if response < FRAMING_MIN_RESPONSE:
        return "測れない", (f"相関が弱く dx/dy が当てにならない "
                            f"(response={response:.3f} < {FRAMING_MIN_RESPONSE:.2f})")
    bad = []
    if abs(dy) > FRAMING_MAX_DY:
        bad.append(f"|dy|={abs(dy):.1f} > {FRAMING_MAX_DY:.0f}")
    if abs(dx) > FRAMING_MAX_DX:
        bad.append(f"|dx|={abs(dx):.1f} > {FRAMING_MAX_DX:.0f}")
    if bad:
        return "ずれ", " / ".join(bad)
    return "OK", (f"|dy|<={FRAMING_MAX_DY:.0f} / |dx|<={FRAMING_MAX_DX:.0f} / "
                  f"response>={FRAMING_MIN_RESPONSE:.2f}")


def open_top_camera():
    cams = camera_resolver.resolve_cameras(require_top=True, require_side=False, verbose=False)
    cap = cv2.VideoCapture(cams["top"])
    if not cap.isOpened():
        sys.exit("[ERROR] 上カメラを開けない")
    t0 = time.time()
    while time.time() - t0 < 8.0:           # 連係カメラは映像が来るまで数秒かかることがある
        ok, f = cap.read()
        if ok and f is not None:
            return cap
        time.sleep(0.2)
    cap.release()
    sys.exit("[ERROR] 上カメラから映像が来ない. iPhone のロック解除・連係カメラの接続を確認 "
             "(Photo Booth の『カメラ』メニューに iPhone が出るか)")


def _put(img, text, org, scale, color, thick):
    """黒で縁取ってから描く. 机の白と マットの黒 のどちらに重なっても読めるように."""
    cv2.putText(img, text, org, cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), thick + 3, cv2.LINE_AA)
    cv2.putText(img, text, org, cv2.FONT_HERSHEY_SIMPLEX, scale, color, thick, cv2.LINE_AA)


def draw_overlay(small_bgr, ref_bgr, dx, dy, resp, alpha, scale, ref_label=""):
    ref = ref_bgr
    if ref.shape[:2] != small_bgr.shape[:2]:
        ref = cv2.resize(ref, (small_bgr.shape[1], small_bgr.shape[0]), interpolation=cv2.INTER_AREA)
    view = cv2.addWeighted(small_bgr, 1.0 - alpha, ref, alpha, 0)
    big = cv2.resize(view, None, fx=scale, fy=scale, interpolation=cv2.INTER_NEAREST)
    verdict, _why = framing_verdict(dx, dy, resp)
    # cv2.putText は日本語を出せないので画面はローマ字. 端末には日本語で出す.
    label = {"OK": "OK", "ずれ": "ZURE", "測れない": "NO MATCH"}[verdict]
    color = (0, 200, 0) if verdict == "OK" else (0, 0, 255)
    _put(big, label, (16, 56), 1.8, color, 4)
    _put(big, f"dy={dy:+6.1f}  dx={dx:+6.1f}  (OK: both <= {FRAMING_MAX_DY:.0f})", (16, 92), 0.6, color, 2)
    _put(big, f"resp={resp:.3f}  (OK: >= {FRAMING_MIN_RESPONSE:.2f})   q: quit", (16, 116),
         0.6, (230, 230, 230), 1)
    if ref_label:        # どのお手本を重ねているか (ブロックの位置を合わせるときの目印)
        _put(big, f"ref: {ref_label}", (16, big.shape[0] - 14), 0.6, (230, 230, 230), 1)
    return big


def main():
    ap = argparse.ArgumentParser(description="上カメラの写り方をお手本に合わせる (アームは動かさない・表示だけ)")
    ap.add_argument("--raw-dir", required=True,
                    help="基準画像に使う学習エピソード (必須. 既定は最初の 1 本の 1 枚目を基準にする)")
    ap.add_argument("--episode", default=None,
                    help="基準にするエピソード名 (例 run_0001_20261019_192343. "
                         "1 本だけに当てはまる部分一致 192343 でもよい). "
                         "未指定なら最初のエピソード")
    ap.add_argument("--alpha", type=float, default=0.5, help="基準画像の重ね具合 (0=ライブのみ / 1=基準のみ)")
    ap.add_argument("--scale", type=float, default=3.0, help="画面の拡大倍率")
    args = ap.parse_args()

    ref_bgr, ref_path = load_reference(Path(args.raw_dir), args.episode)
    ep_dir = ref_path.parent.parent
    j0 = grasp_j0(ep_dir)
    pos = "つかんだ瞬間が不明" if j0 is None else f"つかんだ瞬間の j0 = {j0:+.1f}°"
    ref_label = f"{ep_dir.name}" + ("" if j0 is None else f"  j0={j0:+.1f}")
    print(f"[REF   ] 基準画像: {ref_path}")
    print(f"[REF   ] お手本 {ep_dir.name} ({pos})"
          + ("" if args.episode else "  ← --episode 未指定なので最初のエピソード"))
    print("[REF   ] 重なった画像のブロックに実物のブロックを合わせると, お手本と同じ位置に置ける")
    print(f"[ALIGN ] |dy| <= {FRAMING_MAX_DY:.0f} なら OK. アームは HOME に置いた状態で合わせる. "
          "q または Esc で終了 (ファイルには何も書かない)")
    cap = open_top_camera()
    last = None
    try:
        while True:
            for _ in range(2):            # バッファに溜まった古いフレームを捨てる
                cap.grab()
            ok, frame = cap.read()
            if not ok or frame is None:
                if cv2.waitKey(100) & 0xFF in (ord("q"), 27):
                    break
                continue
            small = to_small_bgr(frame)
            dx, dy, resp = measure_framing(ref_bgr, small)
            cv2.imshow(WINDOW, draw_overlay(small, ref_bgr, dx, dy, resp,
                                            args.alpha, args.scale, ref_label))
            if last is None or abs(dy - last) > 1.0:      # 端末は値が動いたときだけ出す
                verdict, why = framing_verdict(dx, dy, resp)
                print(f"[ALIGN ] dy={dy:+6.1f}  dx={dx:+6.1f}  response={resp:.3f}  "
                      f"→ {verdict} ({why})")
                last = dy
            if cv2.waitKey(30) & 0xFF in (ord("q"), 27):
                break
    except KeyboardInterrupt:
        print("\n[STOP  ] Ctrl-C")
    finally:
        cap.release()
        cv2.destroyAllWindows()
        print("[DONE  ] 終了 (何も書いていません)")


if __name__ == "__main__":
    main()
