import os,json,math,sys,subprocess,time
from pathlib import Path
import numpy as np
import moderngl
from PIL import Image,ImageDraw,ImageFont
from matplotlib.textpath import TextPath
from matplotlib.font_manager import FontProperties
from matplotlib.path import Path as MPath
from scipy.spatial import Delaunay
from shapely.geometry import Polygon
from shapely.ops import triangulate
from shapely import constrained_delaunay_triangles
C=Path(__file__).parent/'cache'; OUT=Path(__file__).parent/'output';OUT.mkdir(exist_ok=True)
S=json.loads((C/'scene.json').read_text());FPS=S['fps']; W=int(os.environ.get('FILM_WIDTH','1280'));H=W*9//16
ctx=moderngl.create_standalone_context(backend='egl');fbo=ctx.simple_framebuffer((W,H),components=3);fbo.use()
vs='''#version 330
in vec3 pos;in vec3 normal;in vec4 ids;in vec4 weights;uniform mat4 bones[100];uniform mat4 vp;uniform mat4 obj;uniform int skinned;out vec3 P;out vec3 N;
void main(){mat4 m=obj;if(skinned==1)m=bones[int(ids.x)]*weights.x+bones[int(ids.y)]*weights.y+bones[int(ids.z)]*weights.z+bones[int(ids.w)]*weights.w;vec4 p=m*vec4(pos,1);P=p.xyz;N=normalize(mat3(m)*normal);gl_Position=vp*p;}
'''
fs='''#version 330
in vec3 P;in vec3 N;out vec4 frag;uniform vec3 color;uniform vec3 eye;uniform float reveal;uniform float alpha;uniform int mode;uniform vec3 center;
void main(){if(P.y>reveal)discard;vec3 n=normalize(N);if(!gl_FrontFacing)n=-n;vec3 v=normalize(eye-P);vec3 l=normalize(vec3(-2,4,4));float diff=max(dot(n,l),0);float rim=pow(1-max(dot(n,v),0),3.4);float spec=pow(max(dot(n,normalize(l+v)),0),60);vec3 c=color*(.23+diff*.70)+vec3(.70,.82,.88)*rim*.32+spec*.26;
if(mode==1){float d=length((P-center).xz);float grid=(1-smoothstep(.009,.018,abs(fract(P.x*.5+.5)-.5)))+(1-smoothstep(.009,.018,abs(fract(P.z*.5+.5)-.5)));c=vec3(.00012,.00018,.00026)+vec3(.022,.029,.037)*exp(-d*.32)+grid*.0019*exp(-d*.32);float shadow=exp(-dot((P-center).xz*vec2(1.5,2.5),(P-center).xz*vec2(1.5,2.5)));c*=1-shadow*.65;}
if(mode==2)c=color*(.5+diff*.7)+vec3(.85,.9,.95)*spec*.8+rim*.12;
if(mode==3)c=color;
float fog=exp(-max(length(eye-P)-7,0)*.055);c=mix(vec3(.00012,.00018,.00026),c,fog);c=pow(max(c,vec3(0)),vec3(1/2.2));frag=vec4(c,alpha);}
'''
prog=ctx.program(vertex_shader=vs,fragment_shader=fs)
identity=np.eye(4,dtype='f4');prog['obj'].write(identity.tobytes());prog['reveal']=100.;prog['alpha']=1.;prog['mode']=0
meshes=[]
for m in S['meshes']:
 data=np.fromfile(C/(m['name']+'.bin'),dtype='f4').reshape(-1,14);vbo=ctx.buffer(data.tobytes());vao=ctx.vertex_array(prog,[(vbo,'3f 3f 4f 4f','pos','normal','ids','weights')]);meshes.append((m,data,vao))
segments=[]
for s in S['segments']:
 a=np.memmap(C/(s['segmentKey']+'.mat'),dtype='f4',mode='r',shape=(s['frames'],S['matrixStride']));d=json.loads((C/(s['segmentKey']+'.json')).read_text());segments.append((s,a,np.array(d['centers']),np.array(d['guides'])))
def norm(v):return v/max(np.linalg.norm(v),1e-9)
def view(eye,target):
 f=norm(target-eye);r=norm(np.cross(f,[0,1,0]));u=np.cross(r,f);m=np.eye(4);m[0,:3]=r;m[1,:3]=u;m[2,:3]=-f;m[:3,3]=-m[:3,:3]@eye;return m
def persp(fov=32):
 a=1/math.tan(math.radians(fov)/2);return np.array([[a/(W/H),0,0,0],[0,a,0,0],[0,0,-100.1/99.9,-20/99.9],[0,0,-1,0]])
def setcam(eye,target,fov=32):
 vp=persp(fov)@view(np.array(eye),np.array(target));prog['vp'].write(vp.astype('f4').T.tobytes());prog['eye'].value=tuple(eye);return vp

def geom(vertices):
 a=np.zeros((len(vertices),14),dtype='f4');a[:,:6]=vertices;v=ctx.buffer(a.tobytes());return ctx.vertex_array(prog,[(v,'3f 3f 4f 4f','pos','normal','ids','weights')])
floor=geom([[-100,0,-100,0,1,0],[100,0,-100,0,1,0],[100,0,100,0,1,0],[-100,0,-100,0,1,0],[100,0,100,0,1,0],[-100,0,100,0,1,0]])
font='/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf';bold='/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf'
def textgeom(text,width,y=0,depth=.065):
 tp=TextPath((0,0),text,size=1,prop=FontProperties(fname=bold));polys=tp.to_polygons();bb=tp.get_extents();scale=width/bb.width;verts=[]
 # Build nested glyph contours and triangulate only covered triangles.
 shapes=[Polygon(q) for q in polys if len(q)>3];solid=None
 for shape in sorted(shapes,key=lambda p:-p.area):solid=shape if solid is None else solid.symmetric_difference(shape)
 for tri in constrained_delaunay_triangles(solid).geoms:
  if not solid.covers(tri):continue
  q=np.asarray(tri.exterior.coords)[:3]
  for z,n,order in [(depth,[0,0,1],[0,1,2]),(0,[0,0,-1],[2,1,0])]:
   for j in order:verts.append([(q[j,0]-bb.x0-bb.width/2)*scale,q[j,1]*scale+y,z,*n])
 for poly in polys:
  for a,b in zip(poly[:-1],poly[1:]):
   n=norm(np.array([b[1]-a[1],a[0]-b[0],0]));a=(a-[bb.x0+bb.width/2,0])*scale+[0,y];b=(b-[bb.x0+bb.width/2,0])*scale+[0,y]
   for q in [[*a,0],[*b,0],[*b,depth],[*a,0],[*b,depth],[*a,depth]]:verts.append([*q,*n])
 return geom(verts),np.array(verts)[:,:3]
title,titlepoints=textgeom('NEX',2.65,1.0,.11);sub,_=textgeom('DYNAMIC WARMUP',2.3,.75,.014);demotext,_=textgeom('DEMO',.47,.52,.013)
# Particle targets are sampled from the actual skinned Nex mesh, not a proxy.
rng=np.random.default_rng(424);mm,data,_=meshes[0];sample=data[rng.choice(len(data),1600,replace=False)];bm=segments[0][1][0,mm['offset']:mm['offset']+mm['bones']*16].reshape(-1,4,4).transpose(0,2,1);pp=np.column_stack([sample[:,:3],np.ones(len(sample))]);target=np.zeros((len(sample),4))
for j in range(4):target+=np.einsum('nij,nj->ni',bm[sample[:,6+j].astype(int)],pp)*sample[:,10+j,None]
target=target[:,:3];origin=titlepoints[rng.choice(len(titlepoints),len(target),replace=True)];origin+=rng.normal(0,.009,origin.shape)
p_prog=ctx.program(vertex_shader='''#version 330
in vec3 p;uniform mat4 vp;uniform float size;void main(){gl_Position=vp*vec4(p,1);gl_PointSize=size;}''',fragment_shader='''#version 330
out vec4 c;uniform vec4 tint;void main(){float d=length(gl_PointCoord-.5);if(d>.5)discard;c=vec4(tint.rgb,tint.a*(1-smoothstep(.1,.5,d)));}''')
pbuf=ctx.buffer(reserve=1600*3*4);pvao=ctx.vertex_array(p_prog,[(pbuf,'3f','p')]);lbuf=ctx.buffer(reserve=1000*3*4);lvao=ctx.vertex_array(p_prog,[(lbuf,'3f','p')]);
def points(arr,vp,color=(.5,.72,.8,.6),size=2,lines=False):
 arr=np.asarray(arr,dtype='f4');buf=lbuf if lines else pbuf;vao=lvao if lines else pvao;buf.write(arr.tobytes());p_prog['vp'].write(vp.astype('f4').T.tobytes());p_prog['size']=size;p_prog['tint']=color;vao.render(moderngl.LINES if lines else moderngl.POINTS,vertices=len(arr))
def smooth(x):x=np.clip(x,0,1);return x*x*(3-2*x)
def model_draw(seg,fi,reveal=100,alpha=1,wire=False,reflect=False):
 s,mat,centers,guides=seg;prog['skinned']=1;prog['mode']=0;prog['reveal']=float(reveal);prog['alpha']=alpha;ctx.wireframe=wire
 for m,data,vao in meshes:
  a=mat[fi,m['offset']:m['offset']+m['bones']*16].copy().reshape(-1,4,4)
  if reflect:a[:,:,1]*=-1
  pad=np.tile(np.eye(4,dtype='f4'),(100,1,1));pad[:m['bones']]=a;prog['bones'].write(pad.tobytes())
  col={'Body':(.25,.63,.81),'Shirt':(.87,.89,.92),'Shorts':(.006,.018,.041),'Shoes':(.004,.008,.016),'NeckLace':(.65,.70,.76),'Bracelet':(.4,.47,.54)}[m['name']]
  if wire:col=(.20,.47,.56)
  prog['color']=col
  if m['name']=='Body':
   for g in m['groups']:
    if 'Eyelash' not in m['materials'][g['materialIndex']]:vao.render(vertices=g['count'],first=g['start'])
  else:vao.render()
 ctx.wireframe=False;prog['alpha']=1.;prog['reveal']=100.

def drawtitle(alpha=1,scale=1,dx=0,dy=0):
 prog['skinned']=0;prog['mode']=2;prog['alpha']=alpha;M=np.eye(4,dtype='f4');M[:3,:3]*=scale;M[:3,3]=[dx,dy,0];prog['obj'].write(M.T.tobytes());prog['color']=(.48,.56,.64);title.render();prog['color']=(.55,.64,.69);sub.render();demotext.render();prog['obj'].write(identity.tobytes());prog['alpha']=1.

# Sparse title card and movement labels are composited after physical rendering.
fontcache={}
def ft(size,b=False):
 k=(size,b)
 if k not in fontcache:fontcache[k]=ImageFont.truetype(bold if b else font,size)
 return fontcache[k]
def overlay(im,t,s=None,local=0):
 d=ImageDraw.Draw(im);scale=W/1280;X=lambda v:int(v*scale)
 d.rectangle((0,0,W,X(29)),fill=(2,4,7));d.rectangle((0,H-X(29),W,H),fill=(2,4,7))
 if 13<=t<S['duration']-7 and s:
  d.text((X(56),X(57)),'NEX',font=ft(X(23),True),fill=(208,218,225));d.text((X(113),X(64)),'/  DYNAMIC WARMUP',font=ft(X(10)),fill=(117,139,151))
  d.text((X(58),H-X(121)),f"{s['order']:02d} / 17",font=ft(X(12)),fill=(106,158,172))
  d.text((X(56),H-X(98)),s['title'].upper(),font=ft(X(24),True),fill=(224,230,232))
  x=X(58);yy=H-X(49);d.line((x,yy,W-X(58),yy),fill=(35,46,53),width=1);d.line((x,yy,int(x+(W-2*x)*(t-13)/(S['duration']-20)),yy),fill=(113,167,181),width=2)
  cycle=min(s['loops'],int(local/s['duration'])+1);label=f"CYCLE {cycle:02d}  /  {s['loops']:02d}";d.text((W-X(190),X(63)),label,font=ft(X(11)),fill=(114,140,154))
 if t>=S['duration']-7:
  d.text((X(70),X(432)),'MOVEMENT. RECONSTRUCTED.',font=ft(X(12)),fill=(130,167,179))
 # Discrete review status belongs on the deliverable; not a performance claim.
 return im

def render(t):
 ctx.enable(moderngl.DEPTH_TEST|moderngl.BLEND|moderngl.PROGRAM_POINT_SIZE);ctx.blend_func=moderngl.SRC_ALPHA,moderngl.ONE_MINUS_SRC_ALPHA;ctx.depth_func='<=';ctx.clear(.016,.02,.027,1,depth=1)
 s=None;local=0;intro=t<13;outro=t>=S['duration']-7
 if intro or outro:
  seg=segments[0];fi=0;c=seg[2][fi].copy();c[1]=.84
  if intro:
   angle=math.radians(-23+36*smooth(t/6));eye=np.array([math.sin(angle)*5.1,1.55,math.cos(angle)*5.1]);targetCam=np.array([0,1.03,0]);vp=setcam(eye,targetCam,34)
  else:
   end=t-(S['duration']-7);eye=np.array([2.35-.25*smooth(end/7),1.52,4.65]);targetCam=c+np.array([-.6,.04,0]);vp=setcam(eye,targetCam,31)
 else:
  k=next(i for i,(s,*_) in enumerate(segments) if t<s['end']);seg=segments[k];s=seg[0];local=t-s['start'];fi=min(s['frames']-1,int(local*FPS));c=seg[2][fi].copy();c[1]=.86
  # Center on pelvis x/z with a symmetric low-pass window, retaining actual root travel.
  i0=max(0,fi-10);i1=min(s['frames'],fi+11);c[[0,2]]=seg[2][i0:i1].mean(axis=0)[[0,2]]
  # Legible whole-body lenses; insert at a cycle boundary when repetitions permit.
  angles=[30,145,85,25,88,68,83,5,75,28,65,10,25,89,88,62,83]
  angle=angles[k];second=local>=s['duration'] and s['loops']>1
  if second:angle=[75,85,25,45,30,85,30,30,30,50,88,30,145,76,55,90,30][k]
  if k in [3,9,10]:angle+=12*(smooth(local/s['clipSec'])-.5)
  rad=math.radians(angle);dist=4.35;eh=1.48
  if k in [2,11,12,13,15,16]:eh=1.03;dist=4.5
  if k in [13,14,15,16] and s['duration']*2<=local<s['duration']*3:
   c[1]=.48;dist=2.65;eh=.70
  eye=c+np.array([math.sin(rad)*dist,eh-c[1],math.cos(rad)*dist]);vp=setcam(eye,c,33)
 prog['center']=tuple(c)
 if not intro or t>10.5:model_draw(seg,fi,alpha=.13,reflect=True)
 prog['skinned']=0;prog['mode']=1;prog['alpha']=.93;prog['obj'].write(identity.tobytes());floor.render();prog['alpha']=1.
 # Distant light traces float behind the athlete, subordinate to the movement.
 traces=[[c[0]-4,.025,c[2]-2],[c[0]+4,.025,c[2]-2],[c[0]-4,.025,c[2]-2.025],[c[0]+4,.025,c[2]-2.025]]
 points(traces,vp,(.11,.17,.20,.4),1,True)
 if intro:
  if t<5.8:drawtitle(float(smooth(t/1.6)*(1-smooth((t-4.2)/1.6))))
  if 4<t<10:
   u=smooth((t-4)/4);pos=origin*(1-u)+target*u;pos[:,2]+=np.sin(np.pi*u)*.45*np.sin(np.arange(len(pos))*.031);points(pos,vp,(.48,.67,.76,float(smooth((t-4)/.7)*(1-smooth((t-8.5)/1.5)))),2.1)
  if 6.7<t<10.7:
   model_draw(seg,0,reveal=2.5*smooth((t-6.7)/1.7),alpha=float(.15*(1-smooth((t-9.6)/1.1))),wire=True)
   guides=seg[3][0];pairs=[];names=S['boneNames']
   for a,b in [('Hip','Head')]+[(f'{side}_{a}',f'{side}_{b}') for side in 'LR' for a,b in [('Thigh','Calf'),('Calf','Foot'),('Upperarm','Forearm'),('Forearm','Hand')]]:
    if a in names and b in names:pairs.extend([guides[names.index(a)],guides[names.index(b)]])
   points(pairs,vp,(.5,.72,.78,.55),1,True);points(guides,vp,(.73,.85,.88,.8),5)
  if t>8.7:model_draw(seg,0,reveal=2.4*smooth((t-8.7)/2.4))
 elif outro:
  model_draw(seg,0)
  prog['skinned']=0;prog['mode']=2;M=np.eye(4,dtype='f4');M[:3,:3]*=.48;M[:3,3]=[-1.62,.60,.15];prog['obj'].write(M.T.tobytes());prog['color']=(.50,.59,.66);title.render();sub.render();prog['obj'].write(identity.tobytes())
 else:
  model_draw(seg,fi)
  # Readback guides shown briefly, derived from actual joints; no invented force metrics.
  if k in [4,8,13,14,16] and .6<local<1.6:
   guides=seg[3][fi];points(guides,vp,(.5,.77,.83,.45),3)
  # Brief spatial scan masks exercise changes without interpolating anatomical motion.
  if local<.18 and k>0:
   xx=-2+4*smooth(local/.18);lines=[[c[0]+xx,0,c[2]+.3],[c[0]+xx,2.2,c[2]+.3]];points(lines,vp,(.5,.7,.76,.6),1,True)
 im=Image.frombytes('RGB',(W,H),fbo.read(components=3,alignment=1)).transpose(Image.Transpose.FLIP_TOP_BOTTOM);im=overlay(im,t,s,local)
 if s and s['order']>1 and local<.32:
  prev=render(s['start']-1/FPS);edge=int(W*smooth(local/.32));im.paste(prev.crop((edge,0,W,H)),(edge,0));d=ImageDraw.Draw(im);d.line((edge,0,edge,H),fill=(127,162,172),width=max(2,W//400))
 if t<1.3 or t>S['duration']-1:
  fade=smooth(t/1.3) if t<1.3 else smooth((S['duration']-t)/1);im=Image.blend(Image.new('RGB',(W,H)),im,float(fade))
 return im
if __name__ == '__main__':
 if '--preview' in sys.argv:
  times=[3,7.8,10.7,14.5]+[s['start']+s['duration']*.25 for s in S['segments']]+[S['duration']-4]
  for i,t in enumerate(times):render(t).save(OUT/f'preview_{i:02d}.jpg')
  thumbs=[]
  for i in range(len(times)):
   im=Image.open(OUT/f'preview_{i:02d}.jpg');im.thumbnail((480,270));thumbs.append(im)
  sheet=Image.new('RGB',(480*4,300*math.ceil(len(thumbs)/4)),(14,18,24));d=ImageDraw.Draw(sheet)
  for i,im in enumerate(thumbs):x=i%4*480;y=i//4*300;sheet.paste(im,(x,y));d.text((x+10,y+275),f'{i:02d} / {times[i]:.1f}s',fill='white')
  sheet.save(OUT/'contact_sheet.jpg');print('preview complete',S['duration'])
 else:
  p=OUT/'NEX_DYNAMIC_WARMUP_V425_REVIEW.mp4';ff=subprocess.Popen(['ffmpeg','-y','-loglevel','error','-f','rawvideo','-pix_fmt','rgb24','-s',f'{W}x{H}','-r',str(FPS),'-i','-','-an','-c:v','libx264','-preset','fast','-crf','18','-pix_fmt','yuv420p','-movflags','+faststart',str(p)],stdin=subprocess.PIPE)
  start=time.time()
  for f in range(math.ceil(S['duration']*FPS)):
   ff.stdin.write(render(f/FPS).tobytes())
   if f%(FPS*5)==0:print(f'{f/FPS:.0f}s / {S["duration"]:.0f}s, elapsed {time.time()-start:.0f}s',flush=True)
  ff.stdin.close();assert ff.wait()==0;print(p)