"""Render exact saved Pose observations. Color bands are display-only."""
import csv
import hashlib
import json
import os
import pathlib
import time
import uuid
from collections import Counter
from fractions import Fraction
import av
import numpy as np
from PIL import Image,ImageDraw,ImageFont,ImageOps
from contract import HIGH,LOW,MUTED,EDGES,GROUPS,KOREAN,observation,palette

CLIP=os.environ['CLIP'];CACHE=pathlib.Path('/cache');ROOT=pathlib.Path('/eval')
SPEC=json.loads(pathlib.Path('/specs.json').read_text())[CLIP]
PREVIOUS=pathlib.Path('/previous')/CLIP/'data'
CHECK=json.loads((PREVIOUS/'cache_validation.json').read_text());assert CHECK['status']=='PASS'
CACHE_BYTES=(CACHE/'inputs/pose.json').read_bytes();assert hashlib.sha256(CACHE_BYTES).hexdigest()==CHECK['cache_sha256']
ALL=json.loads(CACHE_BYTES);ROWS=[ALL[i] for i in CHECK['selected_samples']]
COLORS={int(k):tuple(v) for k,v in json.loads((PREVIOUS/'identity_palette.json').read_text()).items()}
DEST=pathlib.Path('/download')/CLIP;DATA=DEST/'data';DATA.mkdir(parents=True,exist_ok=True)
OUT=DEST/'Pose_Input_Compare_b025_b020_j010_j005.mp4';REPORT=DATA/'video_validation.json'
if REPORT.exists() and OUT.exists():
    saved=json.loads(REPORT.read_text());assert saved['status']=='PASS' and saved['sha256']==hashlib.sha256(OUT.read_bytes()).hexdigest()
    print('REUSED_COMPLETE '+CLIP);raise SystemExit(0)
assert not OUT.exists(),'Existing final output is protected'
W,H=3200,1800;IW,IH=1856,1044;IX,IY=20,150
FONT={n:ImageFont.truetype('/font.ttc',n,index=1) for n in (22,24,26,28,30,32,34,40)}
LIMIT=(0,0,W,H)


def intersects(a,b):return a[0]<b[2] and b[0]<a[2] and a[1]<b[3] and b[1]<a[3]


def text(draw,rects,xy,value,size=28,color='#edf1f5',limit=LIMIT):
    value=str(value);box=draw.textbbox(xy,value,font=FONT[size])
    assert box[0]>=limit[0] and box[1]>=limit[1] and box[2]<=limit[2] and box[3]<=limit[3],('text overflow',value,box)
    assert not any(intersects(box,b) for b in rects),('text overlap',value,box)
    rects.append(box);draw.text(xy,value,font=FONT[size],fill=color)


def accent(draw,p,color):
    x,y=p;draw.ellipse((x-8,y-8,x+8,y+8),outline=tuple(color),width=2)


def overlay(image,people):
    fit=ImageOps.contain(image,(IW,IH));panel=Image.new('RGB',(IW,IH),'black')
    ox,oy=(IW-fit.width)//2,(IH-fit.height)//2;panel.paste(fit,(ox,oy));d=ImageDraw.Draw(panel)
    sx,sy=fit.width/image.width,fit.height/image.height
    def point(x,y):return (ox+min(fit.width-1,max(0,x*sx)),oy+min(fit.height-1,max(0,y*sy)))
    geometry=[];protected=[]
    for p in people:
        a,b=point(*p['bbox'][:2]),point(*p['bbox'][2:]);points=[point(*k[:2]) if s in ('HIGH','LOW') else None for k,s in zip(p['joints'],p['joint_states'])]
        geometry.append((p,(*a,*b),points))
        for xy in points:
            if xy:protected.append((xy[0]-9,xy[1]-9,xy[0]+9,xy[1]+9))
        for (j,k),s in zip(EDGES,p['edge_states']):
            if s:
                x,y=points[j],points[k];protected.append((min(x[0],y[0])-4,min(x[1],y[1])-4,max(x[0],y[0])+4,max(x[1],y[1])+4))
    labels=[];occupied=[];far_labels=0
    for p,b,points in geometry:
        name=p['identity'] if p['track_id'] is not None else '미추적 '+p['identity']
        label=f"{name}  {p['person_score']:.3f}"+('  #'+str(p['detail_slot']) if p['detail_slot'] else '')
        tb=d.textbbox((0,0),label,font=FONT[22]);lw,lh=tb[2]-tb[0]+14,tb[3]-tb[1]+10
        near=[(b[0],b[1]-lh-4),(b[2]-lw,b[1]-lh-4),(b[0],b[3]+4),(b[2]+5,b[1]),(b[0]-lw-5,b[1])]
        grid=[(x,y) for y in range(5,IH-lh,lh+5) for x in range(5,IW-lw,max(100,lw//2))]
        grid.sort(key=lambda p:(p[0]-b[0])**2+(p[1]-b[1])**2)
        chosen=None
        for x,y in near+grid:
            r=(x,y,x+lw,y+lh)
            if x<0 or y<0 or r[2]>=IW or r[3]>=IH:continue
            if any(intersects(r,o) for o in occupied) or any(intersects(r,o) for o in protected):continue
            chosen=r;break
        assert chosen is not None,('No safe label space',CLIP,p['identity'])
        occupied.append(chosen);labels.append((p,label,chosen,tb))
        near_distance=((chosen[0]-b[0])**2+(chosen[1]-b[1])**2)**.5
        if near_distance>50:
            far_labels+=1;d.line((chosen[0]+lw/2,chosen[1]+lh/2,(b[0]+b[2])/2,b[1]),fill=tuple(p['identity_color']),width=1)
        d.rectangle(b,outline=HIGH if p['bbox_state']=='HIGH' else LOW,width=3)
    for p,b,points in geometry:
        for (a,c),state in zip(EDGES,p['edge_states']):
            if state:d.line((*points[a],*points[c]),fill=HIGH if state=='HIGH' else LOW,width=2 if state=='HIGH' else 4)
        for xy,state in zip(points,p['joint_states']):
            if xy is None:continue
            x,y=xy;radius=3 if state=='HIGH' else 4;ring=radius+3
            d.ellipse((x-ring,y-ring,x+ring,y+ring),outline=tuple(p['identity_color']),width=1)
            d.ellipse((x-radius,y-radius,x+radius,y+radius),fill=HIGH if state=='HIGH' else LOW)
    for p,label,r,tb in labels:
        d.rectangle(r,fill='#101820',outline=tuple(p['identity_color']),width=2)
        d.text((r[0]+7-tb[0],r[1]+5-tb[1]),label,font=FONT[22],fill='#ffffff')
    return panel,dict(label_boxes=occupied,label_count=len(labels),far_labels=far_labels,label_overlap=0,label_joint_edge_overlap=0)


def canvas(image,row,people,index):
    shown=[p for p in people if p['visible']];page=Image.new('RGB',(W,H),'#101820');d=ImageDraw.Draw(page);rects=[]
    text(d,rects,(20,16),CLIP+' | bbox / 관절 score 비교',40)
    text(d,rects,(20,82),f"원본 Frame {row['source_frame']} | PTS {row['seconds']:.3f}s | 관측 {index+1}/{len(ROWS)} | 10FPS 표본 · 정상 속도",28)
    panel,stats=overlay(image,shown);page.paste(panel,(IX,IY))
    baseline=sum(p['bbox_state']=='HIGH' for p in shown);added=len(shown)-baseline
    text(d,rects,(20,1210),f"기준 bbox {baseline}개 (>=0.25) + 추가 {added}개 (0.20~0.25 미만) / cache {len(people)}개",30)
    text(d,rects,(20,1253),f"상세 {min(3,len(shown))}개 / 상세 정보 {max(0,len(shown)-3)}개 생략 | 추가 표시는 탐지 정확도 개선을 뜻하지 않습니다.",24,MUTED)
    for i,p in enumerate(shown[:3]):
        x,y=1910,150+i*465;lim=(x,y,3180,y+455)
        d.rectangle(lim,outline='#43515e',width=2);accent(d,(x+25,y+37),p['identity_color'])
        name=p['identity'] if p['track_id'] is not None else '미추적 '+p['identity']
        text(d,rects,(x+45,y+15),f'#{i+1} {name}',34,tuple(p['identity_color']),lim)
        text(d,rects,(x+685,y+20),f"사람 score {p['person_score']:.4f}",30,HIGH if p['bbox_state']=='HIGH' else LOW,lim)
        c=p['counts'];text(d,rects,(x+20,y+70),f"17관절: 높음 {c['HIGH']} / 낮음 {c['LOW']} / 기준미만 {c['BELOW']} / 무효 {c['INVALID']}",26,limit=lim)
        text(d,rects,(x+20,y+114),'관절 묶음',24,MUTED,lim);text(d,rects,(x+250,y+114),'왼쪽 (신체 기준)',24,MUTED,lim);text(d,rects,(x+745,y+114),'오른쪽 (신체 기준)',24,MUTED,lim)
        for j,(name,a,b) in enumerate(GROUPS):
            yy=y+154+j*43;text(d,rects,(x+20,yy),name+(' *' if j>=4 else ''),28,limit=lim)
            for k,xx in ((a,x+250),(b,x+745)):
                state=p['joint_states'][k];v=p['joints'][k][2] if k<len(p['joints']) and len(p['joints'][k])==3 else None
                value='--' if v is None or not np.isfinite(v) else f'{v:.4f}'
                text(d,rects,(xx,yy),KOREAN[state]+'  '+value,28,HIGH if state=='HIGH' else LOW if state=='LOW' else MUTED,lim)
        core=Counter(p['core_states'].values());extra=Counter(p['extra_states'].values())
        text(d,rects,(x+20,y+416),f"중점8 높음/낮음/미달/무효 {core['HIGH']}/{core['LOW']}/{core['BELOW']}/{core['INVALID']} | 보조4 {extra['HIGH']}/{extra['LOW']}/{extra['BELOW']}/{extra['INVALID']}",22,MUTED,lim)
    if not shown:text(d,rects,(1930,190),'cache 객체 없음' if not people else '표시 기준 충족 객체 없음',32)
    for j,line in enumerate([
        'bbox: 초록 >=0.25 / 빨강 0.20~0.25 미만 / 그 미만 미표기',
        '관절: 초록 >=0.10 / 굵은 빨강 0.05~0.10 미만 (참고)',
        '얇은 관절 링과 ID 테두리 = 객체 구분색',
        'T = 실제 cache ID / 미추적 q = 해당 Frame의 출력 칸',
        '같은 q 또는 색의 재등장은 추적·동일인 확인이 아님',
        '연결선은 양 끝의 낮은 등급 / 미달·무효 관절은 연결 안 함',
        '좌표 무효와 낮은 score 구분 / 보간·새 Tracking 없음',
        '표시용 기준만 적용 / 싸움·쓰러짐·정상 행동 판정 없음',
    ]):text(d,rects,(20,1300+j*52),line,28,MUTED)
    text(d,rects,(1930,1575),'상세표는 원래 출력순 첫3객체입니다.',26,MUTED)
    text(d,rects,(1930,1630),'생략 객체의 관절 상태·원래 score도 data에 보존합니다.',26,MUTED)
    text(d,rects,(1930,1685),'관절0.05 미만 미표기 / 전체 관절 필수 조건 아님',26,MUTED)
    return page,stats


started=time.monotonic();all_decisions=[observation(r,COLORS) for r in ROWS]
samples={0,len(ROWS)-1,max(range(len(ROWS)),key=lambda i:sum(p['visible'] for p in all_decisions[i])),
         max(range(len(ROWS)),key=lambda i:sum(p['counts']['LOW'] for p in all_decisions[i] if p['visible']))}
draft=DEST/('.render-'+uuid.uuid4().hex+'.mp4')
video=av.open(str(draft),'w',options={'movflags':'+faststart'});stream=video.add_stream('libx264',rate=10)
stream.width=W;stream.height=H;stream.pix_fmt='yuv420p';stream.options={'crf':'21','preset':'veryfast'}
stream.codec_context.thread_count=2;stream.time_base=stream.codec_context.time_base=Fraction(1,90000)
mapping=[];input_hashes={};facts=DATA/'observation_states.jsonl';max_people=0;max_omitted=0;far_labels=0
previous_hashes=json.loads((PREVIOUS/'input_rgb_sha256.json').read_text())
with facts.open('w') as f:
    for index,(row,people) in enumerate(zip(ROWS,all_decisions)):
        path=CACHE/'inputs/frames'/f"{row['sample']:04d}.png";input_hashes[path.name]=hashlib.sha256(path.read_bytes()).hexdigest()
        assert input_hashes[path.name]==previous_hashes[path.name]
        with Image.open(path) as raw:image=raw.convert('RGB')
        assert image.size==(row['width'],row['height'])
        page,stats=canvas(image,row,people,index)
        if index in samples:page.save(DATA/f'sample_{index:04d}.jpg',quality=88)
        frame=av.VideoFrame.from_ndarray(np.asarray(page),format='rgb24');frame.pts=round((row['seconds']-ROWS[0]['seconds'])*90000);frame.time_base=Fraction(1,90000)
        for packet in stream.encode(frame):video.mux(packet)
        shown=[p for p in people if p['visible']];max_people=max(max_people,len(shown));max_omitted=max(max_omitted,max(0,len(shown)-3));far_labels+=stats['far_labels']
        mapping.append(dict(output_frame=index+1,sample=row['sample'],source_frame=row['source_frame'],source_seconds=row['seconds'],output_pts_seconds=frame.pts/90000,
                            visible=len(shown),baseline=sum(p['bbox_state']=='HIGH' for p in shown),added=sum(p['bbox_state']=='LOW' for p in shown),detail=min(3,len(shown)),omitted=max(0,len(shown)-3)))
        compact=[]
        for p in people:
            q={k:v for k,v in p.items() if k!='joints'};q['joint_scores']=[k[2] if len(k)==3 else None for k in p['joints']];compact.append(q)
        f.write(json.dumps(dict(sample=row['sample'],source_frame=row['source_frame'],seconds=row['seconds'],persons=compact,labels=stats),ensure_ascii=False,separators=(',',':'))+'\n')
        if index%200==0:print('rendered',CLIP,index+1,flush=True)
for packet in stream.encode():video.mux(packet)
video.close()
with av.open(str(draft)) as inp:
    inp.streams.video[0].codec_context.thread_count=2
    times=[]
    for frame in inp.decode(video=0):
        assert not frame.is_corrupt and (frame.width,frame.height)==(W,H)
        times.append(float(frame.pts*frame.time_base))
assert len(times)==len(ROWS) and all(b>a for a,b in zip(times,times[1:]))
error=max(abs(t-m['output_pts_seconds']) for t,m in zip(times,mapping));assert error<1/90000
assert hashlib.sha256((CACHE/'inputs/pose.json').read_bytes()).hexdigest()==CHECK['cache_sha256']
draft.rename(OUT)
with (DATA/'frame_map.csv').open('w',newline='') as f:
    writer=csv.DictWriter(f,fieldnames=mapping[0]);writer.writeheader();writer.writerows(mapping)
(DATA/'input_rgb_sha256.json').write_text(json.dumps(input_hashes,indent=2))
(DATA/'cache_validation.json').write_text(json.dumps(CHECK,indent=2))
(DATA/'pose_summary.json').write_bytes((CACHE/'reports/pose_summary.json').read_bytes())
(DATA/'identity_palette.json').write_text(json.dumps(COLORS,indent=2))
report=dict(status='PASS',clip=CLIP,frames=len(ROWS),first_source_frame=ROWS[0]['source_frame'],last_source_frame=ROWS[-1]['source_frame'],
    first_seconds=ROWS[0]['seconds'],last_seconds=ROWS[-1]['seconds'],nominal_fps=10,resolution=[W,H],max_pts_error=error,
    all_frame_text_overflow=0,all_frame_text_overlap=0,label_overlap=0,label_joint_edge_overlap=0,
    max_visible=max_people,max_detail_omitted=max_omitted,relocated_labels=far_labels,sample_indices=sorted(samples),
    source_sha256=SPEC['sha256'],cache_sha256=CHECK['cache_sha256'],sha256=hashlib.sha256(OUT.read_bytes()).hexdigest(),
    elapsed_seconds=round(time.monotonic()-started,2),cpu_only=True,engine_or_model_decision=False)
REPORT.write_text(json.dumps(report,indent=2));print(json.dumps(report),flush=True)