"""フィジカルAI研究室の記事サムネイル用に、KIOXIA の実データを小さな JSON に書き出す（読むだけ。元データは変えない）。
使い方: cd /Volumes/KIOXIA/r2r2r-data && python3 <このファイル> <出力.json>
描き方は p2m-site の tools/lab_thumbs/lab_thumb.py。"""
import json, csv, glob, os, sys, re, random
import numpy as np
R = lambda v: round(float(v), 1)

def chunks(d):
    pts = []
    for l in open(d + '/log.jsonl'):
        x = json.loads(l)
        if 'actions' in x and x.get('executed'):
            n = x.get('used_step') or len(x['actions'])
            for a in x['actions'][:n]:
                pts.append([R(a[0]), R(a[1]), R(a[2]), 1 if a[6] > 0.5 else 0, 'm'])
        elif x.get('kind') == 'correction' or ('reached_angles' in x and x.get('reached_angles')):
            if pts and x.get('reached_angles'):
                a = x['reached_angles']; g = pts[-1][3]
                pts.append([R(a[0]), R(a[1]), R(a[2]), g, 'c'])
    return pts

def states(f):
    g = 0; pts = []; last = None
    for r in csv.DictReader(open(f)):
        if r['event_type'] == 'gripper_cmd' and r['gripper_state'] != '': g = int(float(r['gripper_state']))
        if r['event_type'] == 'actual' and r['stage'] in ('model', 'correct'):
            p = [R(r['j0']), R(r['j1']), R(r['j2']), g, 'm' if r['stage'] == 'model' else 'c']
            if p != last: pts.append(p); last = p
    return pts[::2]

def runs(pattern):
    out = []
    for d in sorted(glob.glob('runs/pi05_policy/' + pattern)):
        f = d + '/recording/states.csv'
        p = states(f) if os.path.exists(f) else chunks(d)
        if len(p) > 5: out.append(p)
    return out

def loss(f):
    return [[int(m.group(1)), float(m.group(2))] for m in re.finditer(r'Step (\d+):.*?loss=([0-9.]+)', open(f).read())]

def demos(ds, names=None):
    out = []
    for e in sorted(glob.glob(f'datasets/{ds}/raw_episodes/*/episode.npz')):
        name = e.split('/')[-2]
        if names is not None and name not in names: pass
        a = np.load(e)['state']
        out.append({'name': name, 'p': [[R(s[0]), R(s[1]), R(s[2]), 1 if s[6] > 0.5 else 0, 'h'] for s in a]})
    return out

o = {}
o['LOG345'] = runs('pi05_20261017_1*') + runs('pi05_20261017_20*')
o['LOG348'] = runs('pi05_20261017_22*')
o['LOG358'] = runs('pi05_20261021_08*')
o['LOG361'] = runs('pi05_20261021_1[89]*') + runs('pi05_20261021_2*') + runs('pi05_20261022_*')
o['LOG363'] = runs('pi05_20261024_*')
o['LOG401'] = runs('pi05_20261029_1[0-3]*')
o['LOG403'] = runs('pi05_20261029_1[4-9]*')
o['LOG407'] = runs('pi05_20261030_*')
o['LOG417'] = runs('pi05_20261031_*')
o['LOG422'] = runs('pi05_20261101_*')
o['LOG450'] = runs('pi05_20261106_*')
o['loss'] = {'LOG341': loss('checkpoints/pi05_mecharm_lora_v1/logs/train.log'),
             'LOG348': loss('checkpoints/pi05_mecharm_lora_v1/logs/train.log') + loss('lab/Log06/train_resume.log'),
             'LOG355': loss('lab/Log10/logs/train.log'),
             'LOG358': loss('lab/Log11/pod_trim2_logs/train.log'),
             'LOG417': loss('lab/Log17/train/train_trim_corr_v3.log'),
             'LOG422': loss('lab/Log18/train/train_trim_corr_v4.log'),
             'LOG450': loss('lab/Log22/train/train_trim_corr_v5.log')}
o['LOG337'] = [x['p'] for x in demos('mecharm_redblock_box_scripted_v1')][:2]
o['LOG352'] = [x['p'] for x in demos('mecharm_redblock_box_teach_v1')]
# LOG460: 学習データ30本、浅いつかみの7本に印
shallow = set(re.findall(r'\| (\S+) \| \S+ \| \d+ \| 浅い', open('datasets/mecharm_redblock_box_teach_v1_trim2_corr_v5_deep/README.md').read()))
ep = demos('mecharm_redblock_box_teach_v1_trim2') + demos('mecharm_redblock_box_corr_v5')
o['LOG460'] = [{'shallow': x['name'] in shallow, 'p': x['p']} for x in ep]
# LOG349: 手で教えた動きの再生（drag_v2）
o['LOG349'] = []
for f in sorted(glob.glob('runs/experiments/waypoint_streaming_replay/LOG349_20261018/drag_v2_*.json')):
    d = json.load(open(f)); ev = {e['idx']: e['action'] for e in d['gripper_events']}; g = 0; p = []
    for i, a in enumerate(d['angles']):
        if i in ev: g = 1 if ev[i] == 'close' else 0
        p.append([R(a[0]), R(a[1]), R(a[2]), g, 'h'])
    o['LOG349'].append(p[::2])
# LOG327: LIBERO の成否と GPU メモリ
o['LOG327'] = {'success': [l.strip().endswith('True') for l in open('lab/Log03/logs/client.log') if 'Success:' in l],
               'mem': [int(r[' memory.used [MiB]'].split()[0]) for r in csv.DictReader(open('lab/Log03/logs/gpu.csv'))]}
# LOG446: 長押しの角度の動き（jog_probe）
jp = json.load(open('diag/jog_probe_20261105_153401/result.json'))
o['LOG446'] = {k: v.get('samples') for k, v in jp['tests'].items() if isinstance(v, dict) and v.get('samples')}
# LOG419: ゲームパッドの押した記録（試運転 11/02）
o['LOG419'] = [[json.loads(l).get('t_unix'), json.loads(l).get('button'), json.loads(l).get('joint'), json.loads(l).get('delta')] for l in open('runs/experiments/gamepad_teleop/teleop_20261102_192657_dry/log.jsonl')]
# LOG461: 爪の電流（サーボ7 の番地69・70）
base = 'diag/servo_current_20261107/180051_pad/'; cur = []
for r in csv.DictReader(open(base + 'samples.csv')):
    try:
        v = {x[1]: x[2] for x in json.loads(r['regs']) if x[0] == 7}; cur.append((v[69] + v[70] * 256) & 0x7FFF)
    except Exception: cur.append(None)
o['LOG461'] = cur
# LOG225: FK テーブルの点（爪先の位置）
random.seed(1); pts = []
with open('fk_tables/fk_scan_motion_based_v2.csv') as f:
    next(f)
    for i, l in enumerate(f):
        if i % 280 == 0:
            s = l.split(','); pts.append([round(float(s[0]), 4), round(float(s[2]), 4)])
o['LOG225'] = pts
# LOG277: これまでの実機の通し試行（台帳）の置いた位置
o['LOG277'] = [[R(r['block_place_x'] or 0), R(r['block_place_z'] or 0), r['place_verdict']] for r in csv.DictReader(open('runs/pick_place_ledger_v3.csv')) if r['block_place_x'] and r['block_place_z']]
json.dump(o, open(sys.argv[1], 'w'), separators=(',', ':'))
for k, v in o.items(): print(k, len(v) if hasattr(v, '__len__') else v)