import colorsys
import csv
import hashlib
import json
import math
import os
import pathlib
import shutil
import sys
from fractions import Fraction
import av
import numpy as np
from PIL import Image,ImageDraw,ImageFont

ROOT=pathlib.Path('/eval'); clip=os.environ['CLIP']; name=sys.argv[1]
OUT=pathlib.Path('/download')/clip; DATA=OUT/'data'; DATA.mkdir(exist_ok=True)
camera=clip=='01_camera24'
rows=json.loads((ROOT/'inputs/pose.json').read_text())
results=json.loads((ROOT/'outputs'/f'{name}.json').read_text())
assert json.loads((ROOT/'reports'/f'{name}_summary.json').read_text())['status']=='PASS'
labels={int(r['output_index']):r['class_name'] for r in csv.DictReader(open('/work/reports/classes_'+name+'.csv'))}
ntu=name in ('stgcnpp','posec3d')
main=[23,49,50,51] if ntu else [9,19,34,35,64,71]
aux=[22,26,30,41,52,53,54,56,58] if ntu else [1,3,4,5,6,7,8,10,42,43,45,46,50,58,59,65,66,68,70,73,75,76,78]
offset=1 if name=='yowov2' else 0
model={'stgcnpp':'STGCNpp','posec3d':'PoseC3D','ava_slowonly':'AVA_SlowOnly','yowov2':'YOWOv2'}[name]
fonts={n:ImageFont.truetype('/fonts/NotoSansCJK-Regular.ttc',n,index=1) for n in (22,26,30,36)}
orange='#ffae57'; purple='#c6a1ff'; white='#edf1f5'; gray='#c5c9cd'
ids=sorted({p['track_id'] for r in rows for p in r['persons'] if p['track_id'] is not None})
palette={tid:tuple(round(x*255) for x in colorsys.hsv_to_rgb((j*.61803398875)%1,.65,.98)) for j,tid in enumerate(ids)}
audit=[]

def text(d,xy,s,size=26,color=white,right=3190):
    b=d.textbbox(xy,str(s),font=fonts[size])
    assert b[2]<=right and b[3]<=H,(s,b,W,H,right)
    d.text(xy,str(s),font=fonts[size],fill=color)

def ranked(item):
    scores=item['scores']
    indices=list(range(60)) if ntu else (list(range(1,81)) if name=='ava_slowonly' else list(range(80)))
    assert len(scores)==len(indices) and all(math.isfinite(v) and 0<=v<=1 for v in scores)
    table={idx+offset:(labels[idx],scores[j]) for j,idx in enumerate(indices)}
    return table,sorted(table,key=lambda k:(-table[k][1],k))

for limit in ((9,5) if ntu else (10,5)):
    file=OUT/f'{model}_aux{limit:02d}.mp4'; report=DATA/f'{model}_aux{limit:02d}_validation.json'
    if report.exists():
        assert json.loads(report.read_text())['status']=='PASS'; continue
    assert not file.exists(),'Partial render exists: inspect before replacing'
    count=max(1,max(sum('scores' in it for it in r['items']) for r in results))
    W=3200; CARD=110+34*(3+len(main)+limit)+3*38
    H=1040+math.ceil(count/3)*CARD+50; H+=H%2
    c=av.open(str(file),'w',options={'movflags':'+faststart'}); stream=c.add_stream('libx264',rate=10)
    stream.width=W;stream.height=H;stream.pix_fmt='yuv420p';stream.options={'preset':'veryfast','crf':'21'}
    stream.codec_context.thread_count=2;stream.time_base=Fraction(1,90000);stream.codec_context.time_base=Fraction(1,90000)
    first=results[0]['seconds']; maps=[]; maximum_label_count=0
    sample_indices={0,len(results)//2,len(results)-1,max(range(len(results)),key=lambda i:len(results[i]['items']))}
    for index,r in enumerate(results):
        frame=rows[r['sample']]
        assert frame['source_frame']==r['source_frame'] and frame['seconds']==r['seconds']
        rgb=(pathlib.Path('/old/inputs/frames') if camera else ROOT/'inputs/frames')/f"{r['sample']:04d}.png"
        image=Image.open(rgb).convert('RGB'); assert image.size==(frame['width'],frame['height'])
        canvas=Image.new('RGB',(W,H),'#111820'); canvas.paste(image.resize((1600,900)),(0,110)); d=ImageDraw.Draw(canvas)
        text(d,(20,10),f'{clip} | {model} | AUX {limit}',36)
        text(d,(20,60),f"Source Frame {r['source_frame']} | PTS {r['seconds']:.3f}s | sample {r['sample']} | 10FPS",30)
        text(d,(1650,125),'표시: 전체 Top3 + 주 신호 전체 + 보조 신호',30)
        text(d,(1650,180),'주 후보: 주황색 / 보조 후보: 보라색',30)
        text(d,(1650,235),'후보는 미검증 목록이며 싸움 확정 표시가 아닙니다.',26)
        text(d,(1650,290),'같은 클래스가 여러 목록에 중복 표시됩니다.',26)
        text(d,(1650,345),'T = 실제 ID / 미추적(q) = 현재 출력 칸',26)
        text(d,(1650,400),'Y = 이번 Frame의 출력 순서 (Tracking ID 아님)',26)
        text(d,(1650,455),'score는 원래 모델 점수; 모델 간 같은 확률이 아닙니다.',26)
        text(d,(1650,510),'objectness는 행동 점수와 별개입니다.',26)
        window=r.get('window')
        if window:
            text(d,(1650,580),f"입력 {len(window['samples'])}관측: {window['start_seconds']:.3f} ~ {window['end_seconds']:.3f}s",30)
            text(d,(1650,635),f"미래 Frame 사용: {'있음' if window['end_seconds']>r['seconds'] else '없음'}",26)
        else: text(d,(1650,580),'입력 구간 부족: 결과를 채우지 않음',30)
        if name=='stgcnpp': text(d,(1650,700),'쌍 조건: bbox 가장자리 간격 / 작은 높이 ≤ 0.25',26)
        if name=='ava_slowonly': text(d,(1650,700),'AVA adapter: native 수치 동등성 미검증',26)
        text(d,(1650,755),'정탐·오탐·모델 채택은 사용자 검토 대상입니다.',26)
        sx=1600/frame['width']; sy=900/frame['height']; occupied=[]
        people=r['items'] if name=='yowov2' else frame['persons']
        for j,p in enumerate(people):
            if 'bbox' not in p:continue
            tid=p.get('track_id');label=p.get('display_id') or (f'T{tid}' if tid is not None else f"미추적(q{p['query']})")
            color=palette.get(tid,gray) if name!='yowov2' else '#64d8ff'
            x1,y1,x2,y2=p['bbox']; box=(max(0,x1*sx),max(110,y1*sy+110),min(1599,x2*sx),min(1009,y2*sy+110))
            if box[2]<box[0] or box[3]<box[1]:continue
            d.rectangle(box,outline=color,width=3)
            tw=d.textlength(label,font=fonts[22])+10
            for step in range(200):
                x=max(0,min(1598-tw,box[0]+(step//20)*150));y=max(112,min(978,box[1]-32-(step%20)*30))
                rect=(x,y,x+tw,y+29)
                if not any(rect[0]<a[2]+2 and rect[2]+2>a[0] and rect[1]<a[3]+2 and rect[3]+2>a[1] for a in occupied):break
            else: raise RuntimeError('ID label layout capacity exceeded')
            occupied.append(rect); d.line((x+tw/2,y+29,box[0],box[1]),fill=color,width=1);d.rectangle(rect,fill='#111820')
            text(d,(x+4,y),label,22,color,right=1600)
        maximum_label_count=max(maximum_label_count,len(occupied))
        items=[it for it in r['items'] if 'scores' in it]
        if not items: text(d,(24,1055),r.get('status','NO_ELIGIBLE_OUTPUT')+' / 표시할 행동 출력 없음 (정상 행동 판정 아님)',30)
        for j,it in enumerate(items):
            x=20+(j%3)*1060;y=1050+(j//3)*CARD
            if name=='stgcnpp': label=' + '.join('T'+str(t) for t in it['ids'])+' / '+it['status']
            elif name=='posec3d':label='화면 전체 / 특정 사람에게 귀속하지 않음'
            elif name=='yowov2':label=f"{it['display_id']} | objectness {it['objectness']:.4f}"
            else: label=f"T{it['track_id']}" if it['track_id'] is not None else f"미추적(q{it['query']})"
            d.rectangle((x-5,y-5,x+1030,y+CARD-20),outline='#43515f',width=1)
            text(d,(x,y),label,26,right=x+1025); y+=52
            table,order=ranked(it)
            groups=[('전체 TOP3',order[:3]),('주 신호 후보',sorted(main,key=lambda k:(-table[k][1],k))),('보조 신호 후보',sorted(aux,key=lambda k:(-table[k][1],k))[:limit])]
            assert len(groups[1][1])==len(main) and len(groups[2][1])==limit
            for title,keys in groups:
                text(d,(x,y),title,26);y+=38
                assert all(table[a][1]>=table[b][1] for a,b in zip(keys,keys[1:]))
                for key in keys:
                    classname,score=table[key]; color=orange if key in main else (purple if key in aux else white)
                    text(d,(x,y),f'{key:02d} {classname}',22,color,right=x+875)
                    value='<0.0001' if 0<score<.0001 else f'{score:.4f}'
                    text(d,(x+895,y),value,22,color,right=x+1025);y+=34
        if index in sample_indices:canvas.save(DATA/f'{model}_aux{limit:02d}_sample{index}.jpg',quality=85)
        vf=av.VideoFrame.from_ndarray(np.asarray(canvas),format='rgb24');vf.pts=round((r['seconds']-first)*90000);vf.time_base=Fraction(1,90000)
        for packet in stream.encode(vf):c.mux(packet)
        maps.append(dict(output_frame=index+1,sample=r['sample'],source_frame=r['source_frame'],source_seconds=r['seconds'],output_pts_seconds=vf.pts/90000))
    for packet in stream.encode():c.mux(packet)
    c.close()
    with av.open(str(file)) as video: times=[float(f.pts*f.time_base) for f in video.decode(video=0)]
    assert len(times)==len(results) and all(abs(t-m['output_pts_seconds'])<.001 for t,m in zip(times,maps))
    with (DATA/f'{model}_aux{limit:02d}_frame_map.csv').open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=maps[0]);w.writeheader();w.writerows(maps)
    check=dict(status='PASS',clip=clip,model=name,aux=limit,frames=len(times),resolution=[W,H],duration_last_pts=times[-1],max_pts_error=max(abs(t-m['output_pts_seconds']) for t,m in zip(times,maps)),max_items=count,max_labels=maximum_label_count,sha256=hashlib.sha256(file.read_bytes()).hexdigest(),color_main=orange,color_aux=purple,main_ids=main,aux_ids=aux,all_class_top=3,overflow=0)
    report.write_text(json.dumps(check,indent=2));print(json.dumps(check),flush=True)
shutil.copyfile(ROOT/'outputs'/f'{name}.json',DATA/f'{name}_raw.json')
shutil.copyfile(ROOT/'reports'/f'{name}_summary.json',DATA/f'{name}_summary.json')
shutil.copyfile('/work/reports/classes_'+name+'.csv',DATA/f'classes_{name}.csv')
