# LOG418 線あり・線なしの比べ（v3：回・分け方・出力先を引数で渡す。中身は v2 と同じ）
# 使い方: python3 make_crops_v3.py <回の名前> <train|test> <出力フォルダ>
# 元: Log12 予備実験 v2：Unity の録画から、配膳ロボットのまわり3m四方を切り出す
# 普段（走っている） = 往復の途中で5秒に0.5m以上動いた時刻。道の記録は走っている時に1m前後ずれるので、
#   記録の位置から1.5m以内で一番近いロボットを録画の中で探し、その位置を中心にする
# 普段（止まっている） = 待機場所で待っている・卓で料理を渡している時刻（ふつうの止まり）
# 短い止まり = 3〜30秒の止まりの、始まりから3秒後／長い止まり = 30秒以上の止まりの、始まりから15秒後
# A = その時刻の1コマ／B = 10秒前・5秒前・その時刻の3コマを白黒にして重ねた画像
import csv,json,os,sys,random,bisect,glob
import numpy as np,cv2
U=os.path.expanduser('~/mnt/KIOXIA/r2r2r-data/service_sim/unity/')
OUT=sys.argv[3].rstrip('/')+'/'
CROP_M=3.0; SIZE=224; random.seed(0)
run=sys.argv[1]; split=sys.argv[2]; d=U+run
der=json.load(open(d+'/conditions.json'))['derived']; OX,OZ=der['outer_size_x_m'],der['outer_size_z_m']
mp4=glob.glob(d+'/recordings/*_whole.mp4')[0]
W,H=(lambda c:(int(c.get(3)),int(c.get(4))))(cv2.VideoCapture(mp4)); s=W/(OX+1.6)
assert abs(H-s*(OZ+1.6))<3,(W,H,s)
tr={}
for r in csv.DictReader(open(d+'/robot_track.csv')): tr.setdefault(r['robot_id'],[]).append((float(r['t_s']),float(r['x']),float(r['z']),r['phase']))
for k in tr: tr[k].sort()
stops=list(csv.DictReader(open(d+'/stops.csv')))
busy={k:[] for k in tr}
for r in stops: busy[r['robot_id']].append((float(r['t_s']),float(r['t_s'])+float(r['duration_s'])))
def inbusy(rid,t,pre=12,post=2): return any(a-pre<=t<=b+post for a,b in busy[rid])
S=[]
for r in stops:
    dur=float(r['duration_s']); t0=float(r['t_s'])
    if dur>=30: S.append(['long',r['robot_id'],t0+15,float(r['x']),float(r['z']),dur,r['cause'],False])
    elif dur>=3: S.append(['short',r['robot_id'],t0+3,float(r['x']),float(r['z']),dur,r['cause'],False])
sh=[x for x in S if x[0]=='short']; random.shuffle(sh); S=[x for x in S if x[0]=='long']+sh[:120]
mv,idle=[],[]
for rid,L in tr.items():
    for (t1,x1,z1,p1),(t2,x2,z2,p2) in zip(L,L[1:]):
        if t2<15 or inbusy(rid,t2): continue
        if p2 in('卓へ','受け渡し口へ') and np.hypot(x2-x1,z2-z1)>=0.5: mv.append(['normal_move',rid,t2,x2,z2,0,'',True])
        if p1==p2 and p2 in('待機','往復おわり','卓で止まる') and np.hypot(x2-x1,z2-z1)<0.05: idle.append(['normal_idle',rid,t2,x2,z2,0,p2,False])
random.shuffle(mv); random.shuffle(idle); S+=mv[:250]+idle[:120]
def blobs(img):
    w=(img>235).all(2).astype(np.uint8); n,lab,st,cen=cv2.connectedComponentsWithStats(w)
    return [cen[i] for i in range(1,n) if 0.06*s*s<st[i,4]<0.4*s*s]
# 1回目：走っている所は、その時刻のコマでロボットを探して中心を決める
cen={}; want={}
for i,x in enumerate(S):
    if x[7]: want.setdefault(int(round(x[2]/0.5)),[]).append(i)
    else: cen[i]=(s*(x[3]+0.8),s*(OZ+0.8-x[4]))
cap=cv2.VideoCapture(mp4); f=0
while want and f<=max(want) and cap.grab():
    if f in want:
        ok,img=cap.retrieve(); bl=blobs(img)
        for i in want[f]:
            px,py=s*(S[i][3]+0.8),s*(OZ+0.8-S[i][4])
            dd=[(np.hypot(b[0]-px,b[1]-py),b) for b in bl]
            if dd and min(dd,key=lambda a:a[0])[0]<1.5*s: cen[i]=tuple(min(dd,key=lambda a:a[0])[1])
    f+=1
# 2回目：切り出す
need={}
for i in cen:
    fi=int(round(S[i][2]/0.5))
    for k,ff in enumerate((fi-20,fi-10,fi)): need.setdefault(ff,[]).append((i,k))
def crop(img,c):
    h=CROP_M*s/2; x0,y0=int(c[0]-h),int(c[1]-h); x1,y1=int(c[0]+h),int(c[1]+h)
    pad=cv2.copyMakeBorder(img,400,400,400,400,cv2.BORDER_CONSTANT,value=(100,100,100))
    return cv2.resize(pad[y0+400:y1+400,x0+400:x1+400],(SIZE,SIZE),interpolation=cv2.INTER_AREA)
parts={}; cap=cv2.VideoCapture(mp4); f=0
while f<=max(need) and cap.grab():
    if f in need:
        ok,img=cap.retrieve()
        for i,k in need[f]: parts[(i,k)]=crop(img,cen[i])
    f+=1
rows=[]; import re; tag=re.sub(r'_cs\d+_fam_p1_','_',run[9:]).split('_s7')[0]
for i,(lab,rid,t,x,z,dur,c,_) in enumerate(S):
    if any((i,k) not in parts for k in range(3)): continue
    a=parts[(i,2)]; b=np.dstack([cv2.cvtColor(parts[(i,k)],cv2.COLOR_BGR2GRAY) for k in range(3)])
    name=f'{tag}_r{rid}_t{t:.1f}.png'; grp='normal' if lab.startswith('normal') else lab
    for kind,im in (('A',a),('B',b)):
        p=f'{OUT}{kind}/{split}/{grp}/'; os.makedirs(p,exist_ok=True); cv2.imwrite(p+name,im)
    rows.append(dict(run=run,split=split,label=lab,robot=rid,t_s=round(t,2),x=x,z=z,cx_px=round(cen[i][0],1),cy_px=round(cen[i][1],1),stop_s=dur,cause=c,file=name))
os.makedirs(OUT+'index',exist_ok=True)
with open(OUT+f'index/index_{run}.csv','w',newline='') as fo:
    w=csv.DictWriter(fo,fieldnames=list(rows[0])); w.writeheader(); w.writerows(rows)
from collections import Counter; print(tag,split,Counter(r['label'] for r in rows),'move dropped',len(mv[:250])-sum(r['label']=='normal_move' for r in rows))