#!/usr/bin/env python3
"""grip_drop_analysis.py — 爪を閉じている間に、爪の電流（gripper_current）が急に 0 付近へ落ちる所を探す（Log30・LOG487）。

入力: runs/pi05_policy/pi05_20261111_*/recording/states.csv・currents.csv・log.jsonl（読むだけ）
見ること（試行ごと）
- 閉じた指令（gripper_cmd の gripper_state=1）の時刻。1 回目を 0 秒にする
- 閉じて 0.8〜1.5 秒の電流の中央値
- 「落ちた」所: 1 つ前の行が 400 を超え, 次の行が 150 を下回った所（爪が閉じている間だけ）
  そのとき, 直前の send_angles（推論が送った目標）から何秒か, 前 1 秒の肩 j1・ひじ j2 の動き（大きい方）
- 締め直し r（log.jsonl の regrip_tighten）の時刻と送った値
- 爪を閉じてから開くまで, 肩 j1 がいちばん上がった所（閉じた所からの差。マイナス＝持ち上げ）
出力: 表を標準出力に, 同じ中身を grip_drop_summary.csv に, 図とサムネイル用の時系列を grip_drop_series.json に（このフォルダ）
"""
import csv
import json
import sys
from pathlib import Path

import numpy as np

R = Path('/Volumes/KIOXIA/r2r2r-data/runs/pi05_policy')
if not R.exists():
    R = Path.home() / 'mnt/KIOXIA/r2r2r-data/runs/pi05_policy'
OUT = Path(__file__).resolve().parent
SERIES_RUNS = ['pi05_20261111_100659', 'pi05_20261111_101557', 'pi05_20261111_102043',
               'pi05_20261111_160851', 'pi05_20261111_161402', 'pi05_20261111_161743']


def load(d):
    rows = list(csv.DictReader(open(d / 'recording/states.csv')))
    cur = [(float(r['ts']), float(r['gripper_current'])) for r in csv.DictReader(open(d / 'recording/currents.csv'))
           if r['gripper_current'] not in ('', None)]
    log = [json.loads(l) for l in open(d / 'log.jsonl')]
    return rows, cur, log


def analyze(d):
    rows, cur, log = load(d)
    gv = next((j.get('grip_value') for j in log if j.get('type') == 'model'), None)
    gc = [(float(r['ts']), r['gripper_state']) for r in rows if r['event_type'] == 'gripper_cmd']
    closes = [t for t, s in gc if s == '1']
    if not closes:
        return None
    t0 = closes[0]
    sends = [float(r['ts']) for r in rows if r['event_type'] == 'send_angles' and r['stage'] != 'home']
    act = [(float(r['ts']), float(r['j1']), float(r['j2'])) for r in rows
           if r['event_type'] == 'actual' and r['j1'] and r['stage'] != 'home']
    c1 = [v for t, v in cur if t0 + 0.8 < t < t0 + 1.5]
    drops = []
    for (ta, va), (tb, vb) in zip(cur, cur[1:]):
        if va > 400 and vb < 150:
            g = [s for t, s in gc if t < tb]
            if not g or g[-1] != '1':
                continue
            last = max([s for s in sends if s < tb], default=None)
            a = [x for x in act if tb - 1.0 < x[0] < tb + 0.3]
            mv = max(abs(a[-1][1] - a[0][1]), abs(a[-1][2] - a[0][2])) if len(a) > 1 else 0.0
            drops.append({'t': round(tb - t0, 1), 'since_send': None if last is None else round(tb - last, 1),
                          'move_1s': round(mv, 1)})
    regrips = [{'t': round(j['t'] - t0, 1), 'value': j.get('grip_value_sent')} for j in log
               if j.get('action') == 'regrip_tighten' and j.get('sent')]
    opens = [t for t, s in gc if s == '0' and t > t0]
    t_open = opens[0] if opens else (act[-1][0] if act else t0)
    seg = [x for x in act if t0 <= x[0] <= t_open]
    j1_close = seg[0][1] if seg else None
    lift = round(min(x[1] for x in seg) - j1_close, 1) if seg else None
    return {'run': d.name, 'grip_value': gv if gv is not None else 20, 'closes': [round(t - t0, 1) for t in closes],
            'cur_1s': round(float(np.median(c1)), 0) if c1 else None, 'drops': drops, 'regrips': regrips,
            'opened': bool(opens), 'j1_lift': lift}


def series(d):
    rows, cur, log = load(d)
    t0 = next(float(r['ts']) for r in rows if r['event_type'] == 'gripper_cmd' and r['gripper_state'] == '1')
    act = [(float(r['ts']) - t0, float(r['j1'])) for r in rows if r['event_type'] == 'actual' and r['j1'] and r['stage'] != 'home']
    gc = [(float(r['ts']) - t0, int(r['gripper_state'])) for r in rows if r['event_type'] == 'gripper_cmd']
    sends = [round(float(r['ts']) - t0, 2) for r in rows if r['event_type'] == 'send_angles' and r['stage'] != 'home']
    end = max(t for t, _ in act)
    cs = [(round(t - t0, 2), v) for t, v in cur if -2 <= t - t0 <= end]
    step = 0.25
    grid = np.arange(-2, end, step)
    def resample(pts):
        out, i = [], 0
        for g in grid:
            while i + 1 < len(pts) and pts[i + 1][0] <= g:
                i += 1
            out.append(round(pts[i][1], 1))
        return out
    return {'run': d.name, 't0': -2, 'step': step, 'current': resample(cs), 'j1': resample(act),
            'grip': gc, 'sends': sends}


def main():
    res = []
    print('run, 閉じる値, 閉じて1秒の電流, 閉じた時刻(秒), 落ちた所(秒・直前の送りから・前1秒の動き), 締め直し r(秒・値), 開いた, 肩j1の持ち上げ(°)')
    for d in sorted(R.glob('pi05_20261111_*')):
        a = analyze(d)
        if a is None:
            continue
        res.append(a)
        dr = ' / '.join(f"{x['t']}s(送り{x['since_send']}s・{x['move_1s']}°)" for x in a['drops']) or '-'
        rg = ' / '.join(f"{x['t']}s({x['value']})" for x in a['regrips']) or '-'
        print(f"{a['run'][-6:]}, {a['grip_value']}, {a['cur_1s']}, {a['closes']}, {dr}, {rg}, {a['opened']}, {a['j1_lift']}")
    with open(OUT / 'grip_drop_summary.csv', 'w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['run', 'grip_value', 'cur_1s', 'closes', 'drops', 'regrips', 'opened', 'j1_lift'])
        for a in res:
            w.writerow([a['run'], a['grip_value'], a['cur_1s'], json.dumps(a['closes']), json.dumps(a['drops'], ensure_ascii=False),
                        json.dumps(a['regrips']), a['opened'], a['j1_lift']])
    ser = {r: series(R / r) for r in SERIES_RUNS}
    json.dump({'summary': res, 'series': ser}, open(OUT / 'grip_drop_series.json', 'w'), ensure_ascii=False)
    print('書き出し', OUT / 'grip_drop_summary.csv', OUT / 'grip_drop_series.json')


if __name__ == '__main__':
    main()