#!/usr/bin/env python3
"""rewind_correct.py — 実機推論の「巻き戻し (b)」と「関節の補正 (c)」の計算部分 (2026-10-29).

アームにも daemon にも触らない (計算だけ). run_policy_client.py と
build_correction_episodes.py の両方から使い, test_rewind_correct.py で確かめる.

  - rewind_waypoints: 実行した区間 (推論1回分 / 補正1回分) で送った目標を逆順にたどり,
    始めの姿勢で終わる中継点の並びを作る
  - parse_correction_key: 補正モードのキー 1 文字を操作に読み替える
  - excluded_windows / step_keep_mask: 学習データから除く時間の窓 (巻き戻しで取り消した区間)
"""
import math

import numpy as np

# 補正の刻み (度). [ / ] で切り替える. 既定は 1°
CORRECTION_STEPS = (0.5, 1.0, 2.0)
CORRECTION_DEFAULT_STEP_INDEX = 1

# 記録の 1 コマごとの状態 (states.csv の stage 列 / top_frame_log.csv の mode 列)
MODE_MODEL = "model"        # モデルの動き (推論結果の実行・推論待ち・確認待ち)
MODE_REWIND = "rewind"      # 巻き戻し中 (補正の取り消しで戻る間も含む)
MODE_CORRECT = "correct"    # 補正中
MODE_HOME = "home"          # HOME へ戻る間
MODE_WAIT = "wait"          # 推論のあとの確認で入力 (Enter / b / c / y/n) を待っている間.
                            # 人がブロックや箱を手で置き直すと手が写るので学習データに使わない (2026-10-29 依頼 27)
TRAIN_MODES = (MODE_MODEL, MODE_CORRECT)   # 学習データに使う状態 (rewind / home / wait は使わない)

# つかんだ区間 (グリッパーを閉じた推論 1 回分) を巻き戻すときの答え (log.jsonl の rewind 行の grasp_undo)
GRASP_PLACED = "placed"     # y = 持っている: つかんだ位置まで下ろして開き (置いて), 始めの姿勢へ戻る
GRASP_DROPPED = "dropped"   # n = 落とした: 今の位置ですぐ開いてから, 始めの姿勢へ戻る

CORRECTION_HELP = (
    "  [補正] 1〜6=関節を選ぶ (j0〜j5) / + か . = 選んだ関節を +1 刻み / - か , = -1 刻み\n"
    "         [ ] = 刻みを 0.5° / 1° / 2° で切り替え / g = グリッパー開閉\n"
    "         Enter = 補正を終えてモデルの推論に戻る / b = 補正を取り消して始めの姿勢へ戻る\n"
    "         ? = この一覧"
)


def split_leg(a, b, max_move):
    """a から b へ, どの関節も 1 回 max_move 以内になるよう等分した点を返す (a は含めず b は含む)."""
    a = np.asarray(a, dtype=float)
    b = np.asarray(b, dtype=float)
    span = float(np.max(np.abs(b - a))) if len(a) else 0.0
    n = max(1, int(math.ceil(span / max_move - 1e-9))) if max_move and max_move > 0 else 1
    return [a + (b - a) * (k / n) for k in range(1, n + 1)]


def rewind_waypoints(sent_targets, start_pose, current=None, max_move=None):
    """実行した区間を逆にたどる中継点の並び.

    sent_targets: その区間で送った目標 (送った順). start_pose: 区間を始めたときの姿勢
    (推論なら measured_state の 6 関節, 補正なら補正を始めたときの実測).
    戻り値は「最後に送った目標の 1 つ手前 → … → 最初に送った目標 → start_pose」.
    (最後に送った目標は今いる場所なので送らない.)
    max_move を渡すと, current (既定は最後に送った目標) から順に, 1 回の移動が
    max_move を超える区間を等分する (apply_limits の移動量制限で途中止まりにしないため).
    """
    start = np.asarray(start_pose, dtype=float)[:6]
    sent = [np.asarray(x, dtype=float)[:6] for x in sent_targets]
    pts = sent[:-1][::-1] + [start]
    if not max_move:
        return pts
    prev = (np.asarray(current, dtype=float)[:6] if current is not None
            else (sent[-1] if sent else start))
    out = []
    for p in pts:
        out.extend(split_leg(prev, p, max_move))
        prev = p
    return out


def parse_correction_key(ch):
    """補正モードの 1 文字を操作に読み替える. 戻り値は (操作, 値) か None (関係ないキー).

    ("select", j) / ("move", +1|-1) / ("step", +1|-1) / ("grip", None) /
    ("accept", None) / ("cancel", None) / ("help", None)
    """
    if ch in ("1", "2", "3", "4", "5", "6"):
        return ("select", int(ch) - 1)
    if ch in ("+", "=", "."):
        return ("move", +1)
    if ch in ("-", "_", ","):
        return ("move", -1)
    if ch == "]":
        return ("step", +1)
    if ch == "[":
        return ("step", -1)
    if ch in ("g", "G"):
        return ("grip", None)
    if ch in ("\n", "\r", ""):          # "" = 入力の終わり (EOF). 姿勢はそのまま補正を終える
        return ("accept", None)
    if ch in ("b", "B"):
        return ("cancel", None)
    if ch == "?":
        return ("help", None)
    return None


def next_step_index(i, d):
    return min(max(i + d, 0), len(CORRECTION_STEPS) - 1)


def excluded_windows(records):
    """log.jsonl の記録から, 学習データに入れない時間の窓 [(t0, t1), ...] を作る.

    - 巻き戻した区間: 区間を始めた時刻 (segment_t[0]) から 巻き戻しが終わった時刻 (t_end) まで.
      区間の実行・その後の推論待ち (間違った姿勢で止まっている間)・巻き戻しの動きをまとめて除く
    - 取り消した補正: 補正を始めた時刻から 始めの姿勢へ戻り終えた時刻まで
    """
    out = []
    for r in records:
        if r.get("type") == "rewind" and r.get("rewound"):
            seg = r.get("segment_t") or [None, None]
            t0 = seg[0] if seg[0] is not None else r.get("t_start")
            t1 = r.get("t_end")
            if t0 is not None and t1 is not None:
                out.append((float(t0), float(t1)))
        elif r.get("type") == "correction" and r.get("action") == "cancel":
            t0, t1 = r.get("t_start"), r.get("t_end")
            if t0 is not None and t1 is not None:
                out.append((float(t0), float(t1)))
    return out


def mode_at(t, mode_ts, modes, default=MODE_MODEL):
    """各時刻の状態 (その時刻までで最後に切り替わった状態)."""
    t = np.asarray(t, dtype=float)
    if len(mode_ts) == 0:
        return np.array([default] * len(t), dtype=object)
    order = np.argsort(mode_ts)
    mts = np.asarray(mode_ts, dtype=float)[order]
    ms = np.asarray(modes, dtype=object)[order]
    k = np.searchsorted(mts, t, side="right") - 1
    return np.where(k >= 0, ms[np.clip(k, 0, None)], default)


def step_keep_mask(t, mode_ts, modes, windows):
    """学習に使うステップ: 状態が model / correct で, 除く窓のどれにも入らないもの."""
    m = mode_at(t, mode_ts, modes)
    keep = np.isin(m, TRAIN_MODES)
    t = np.asarray(t, dtype=float)
    for t0, t1 in windows:
        keep &= ~((t >= t0) & (t <= t1))
    return keep, m