"""Local fresh-STEP prediction sandbox; no Altair calls or solver labels at inference."""

import json,os,sys,time,uuid,threading,subprocess,hashlib,importlib.util

from pathlib import Path

from http.server import ThreadingHTTPServer,SimpleHTTPRequestHandler

import numpy as np

# Optional local CUDA wheel, isolated from the portable CPU environment.
GPU_RUNTIME=Path(__file__).resolve().parents[1]/'.gpu-runtime'
if (GPU_RUNTIME/'READY').exists():sys.path.insert(0,str(GPU_RUNTIME))
import torch

from step_surrogate import StepNet,features,TASKS

ROOT=Path(__file__).resolve().parents[1];STORE=Path(os.environ.get('ACE_STORE',ROOT/'artifacts/prediction_jobs'));STORE.mkdir(parents=True,exist_ok=True)

active=ROOT/'artifacts/active_model.json'
model_dir=os.environ.get('ACE_MODEL_DIR') or (json.loads(active.read_text())['directory'] if active.exists() else 'artifacts/step_model_v002')
CK=ROOT/model_dir/'pointnet.pt'

if not CK.exists():CK=ROOT/'artifacts/step_model_v001/pointnet.pt'

ck=torch.load(CK,map_location='cpu',weights_only=True)

if ck.get('architecture')=='StreamNet':
 from stream_surrogate import StreamNet,features,TASKS
 model=StreamNet(ck['query_dim'])
elif ck.get('architecture')=='DenseNet':

 from dense_surrogate import DenseNet

 model=DenseNet(ck['query_dim'])

else:model=StepNet(ck['query_dim'])

requested_device=os.environ.get('ACE_DEVICE','auto')
DEVICE=torch.device(('cuda' if torch.cuda.is_available() else 'cpu') if requested_device=='auto' else requested_device)
if DEVICE.type=='cuda':torch.backends.cuda.matmul.allow_tf32=False
torch.set_num_threads(3);model.load_state_dict(ck['state_dict']);model.eval();model.to(DEVICE);lock=threading.Lock();jobs={}

warp_model=None;WARP=CK.with_name('warp.pt')

if WARP.exists() and ck.get('architecture') in ('DenseNet','StreamNet'):

 from warp_surrogate import WarpNet

 wck=torch.load(WARP,map_location='cpu',weights_only=True);warp_model=WarpNet(wck['query_dim']);warp_model.load_state_dict(wck['state_dict']);warp_model.eval();warp_model.to(DEVICE)

# The conda environment supplies pythonocc; an existing separate CAD runtime is optional.
CAD_PYTHON=Path(os.environ.get('ACE_CAD_PYTHON',sys.executable))
if not os.environ.get('ACE_CAD_PYTHON') and importlib.util.find_spec('OCC') is None:
 legacy=Path('Z:/Desktop/Omega/bin/hodge-step/python.exe')
 if legacy.exists():CAD_PYTHON=legacy
CAD=CAD_PYTHON.parent
KNOWN={}
for p in (ROOT/'artifacts/step_parts').glob('*/geometry.json'):

 m=json.loads(p.read_text());KNOWN[m['sha256']]=m['id']

for folder in STORE.iterdir():
 if not folder.is_dir():continue
 try:
  j=folder.name;jobfile=folder/'job.json';binaryfile=folder/'prediction.acebin';legacyfile=folder/'prediction.json'
  d=json.loads(jobfile.read_text()) if jobfile.exists() else None
  if d is None or not d.get('last_setup'):
   prediction=None
   if binaryfile.exists():
    import struct
    with binaryfile.open('rb') as f:
     length=struct.unpack('<I',f.read(4))[0];prediction=json.loads(f.read(length))['metadata']
   elif legacyfile.exists():
    if d is not None:
     # Legacy metadata precedes the large vertex/field arrays. Read only that prefix.
     with legacyfile.open(encoding='utf-8') as f:prefix=f.read(65536)
     end=prefix.find('"vertices":')
     if end>0:prediction=json.loads(prefix[:end].rstrip().rstrip(',')+'}')
    else:prediction=json.loads(legacyfile.read_text())
   if prediction:
    if d is None:d={'id':j,'name':prediction['name'],'sha256':prediction['source_sha256'],'known_source':prediction['known_source'],'geometry':prediction['geometry']}
    if prediction.get('gate_points_model_mm'):d['last_setup']={'gates':prediction['gate_points_model_mm'],'gravity':prediction['gravity_direction'],'gate_directions':prediction.get('inlet_direction_model'),'radius':prediction['spec']['geometry']['gate_radius_mm']}
  if d:
   d['status']='ready' if (folder/'volume.npz').exists() else 'needs_mesh';jobs[j]=d
 except (ValueError,KeyError,OSError):pass


def prepare(j):

 folder=STORE/j

 try:

  env=os.environ.copy();env['PATH']=str(CAD/'Library/bin')+os.pathsep+str(CAD/'DLLs')+os.pathsep+env['PATH']

  if not (folder/'geometry.npz').exists():

   res=subprocess.run([str(CAD_PYTHON),str(ROOT/'scripts/prepare_prediction_step.py'),str(folder)],env=env,capture_output=True,text=True,timeout=180)

   (folder/'preprocess.log').write_text(res.stdout+'\n'+res.stderr)

   if res.returncode:raise ValueError((res.stderr or res.stdout)[-1600:])

  mesh=subprocess.run([sys.executable,str(ROOT/'scripts/mesh_prediction_step.py'),str(folder)],capture_output=True,text=True,timeout=180)

  (folder/'mesh.log').write_text(mesh.stdout+'\n'+mesh.stderr)

  if mesh.returncode:raise ValueError('Could not generate a connected display volume: '+mesh.stderr[-800:])

  jobs[j].update(status='ready',geometry=json.loads((folder/'geometry.json').read_text()))
  (folder/'job.json').write_text(json.dumps(jobs[j]))

 except Exception as e:jobs[j].update(status='error',error=str(e))

def model_info():

 report=json.loads(CK.with_name('report.json').read_text()) if CK.with_name('report.json').exists() else {}
 trained=[c for c in report.get('cases',[]) if c.get('split')=='train']
 return {'name':'Ace-Mini','description':'Artificial Casting Engineer — Mini 1.0','version':CK.parent.name,'architecture':ck.get('architecture','StepNet'),'parameters':sum(p.numel() for p in model.parameters())+(sum(p.numel() for p in warp_model.parameters()) if warp_model is not None else 0),'field_parameters':sum(p.numel() for p in model.parameters()),'trajectory_parameters':sum(p.numel() for p in warp_model.parameters()) if warp_model is not None else 0,'training_configurations':len(trained),'training_geometries':len({c['geometry_id'] for c in trained}),'liquid_distance':ck.get('architecture')=='StreamNet','inference_device':str(DEVICE),'inference_hardware':torch.cuda.get_device_name(DEVICE) if DEVICE.type=='cuda' else 'CPU','frontend_version':'2026-09-09-gpu-inference'}



def rotation(opt,dims):

 axis=opt.get('axis','auto');axis=int(np.argmin(dims)) if axis=='auto' else int(axis)

 if axis not in (0,1,2):raise ValueError('Axis must be X, Y or Z.')

 g=np.asarray(opt.get('gravity',(-np.eye(3)[axis]).tolist()),float)

 if g.shape!=(3,) or not np.isfinite(g).all() or np.linalg.norm(g)<1e-9:raise ValueError('Invalid gravity direction.')

 g=g/np.linalg.norm(g);z=-g;x=np.eye(3)[(int(np.argmax(np.abs(z)))+1)%3];x=x-z*np.dot(x,z);x/=np.linalg.norm(x)

 return np.stack([x,np.cross(z,x),z]),g



def config(j,opt):

 a=np.load(STORE/j/'geometry.npz');dims=np.ptp(a['vertices'],axis=0);r,g=rotation(opt,dims)

 v=a['vertices']@r.T;t=-np.r_[(v[:,:2].min(0)+v[:,:2].max(0))/2,v[:,2].min()];v+=t

 volume=np.load(STORE/j/'volume.npz');points=volume['vertices']@r.T+t;surface=a['surface']@r.T+t;cent=a['centers']@r.T+t;n=a['normals']@r.T

 custom=opt.get('gate_points');count=len(custom) if custom is not None else int(opt.get('gates',1))

 if not 1<=count<=8:raise ValueError('Place between one and eight gates. Training covers one or two.')

 if custom is not None:

  q=np.asarray(custom,float)

  if q.shape!=(count,3) or not np.isfinite(q).all():raise ValueError('Invalid gate coordinates.')

  if opt.get('coordinate_space')=='model':q=q@r.T+t

  # Project onto the closest CAD triangle; this preserves surface clicks rather than snapping to centroids.

  import trimesh

  tri=v[a['triangles']];tri=tri[np.linalg.norm(np.cross(tri[:,1]-tri[:,0],tri[:,2]-tri[:,0]),axis=1)>1e-10];gates=[];face_inward=[]

  for point in q:

   closest=trimesh.triangles.closest_point(tri,np.broadcast_to(point,(len(tri),3)))

   distances=np.linalg.norm(closest-point,axis=1);idx=np.argmin(distances)

   if distances[idx]>max(.5,float(dims.max())*.02):raise ValueError('Gate must be placed on the part surface.')

   if any(np.linalg.norm(closest[idx]-other)<.1 for other in gates):raise ValueError('Duplicate gate locations. Place each gate at a distinct surface point.')
   gates.append(closest[idx]);normal=np.cross(tri[idx,1]-tri[idx,0],tri[idx,2]-tri[idx,0]);face_inward.append(-normal/np.linalg.norm(normal))

 else:

  if count>2:raise ValueError('Place gates manually for more than two gates.')

  candidates=np.where((n[:,2]<-.65)&(cent[:,2]<np.quantile(v[:,2],.65))&(a['areas']>1e-8))[0]

  if len(candidates)<count:raise ValueError('No suitable patches. Place gates manually.')

  rng=np.random.default_rng(int(opt.get('seed',174)));idx=rng.choice(candidates,p=a['areas'][candidates]/a['areas'][candidates].sum());gates=[cent[idx]]

  if count==2:gates.append(cent[candidates[np.argmax(np.linalg.norm(cent[candidates]-gates[0],axis=1))]])

 vol=jobs[j]['geometry']['volume_mm3'];radius=float(opt.get('radius') or np.clip(np.sqrt(vol/(500*3*np.pi*count)),3,12))

 if not np.isfinite(radius) or radius<.1 or radius>100:raise ValueError('Gate radius must be between 0.1 and 100 mm.')

 explicit=opt.get('gate_directions')
 if explicit is not None:
  directions=np.asarray(explicit,dtype=float)
  if directions.shape!=(count,3) or not np.isfinite(directions).all() or np.any(np.linalg.norm(directions,axis=1)<1e-9):raise ValueError('Each gate needs one finite, nonzero flow direction.')
  directions=directions/np.linalg.norm(directions,axis=1,keepdims=True)
  if opt.get('coordinate_space')=='model':directions=directions@r.T
 elif custom is not None:directions=np.asarray(face_inward)
 else:
  directions=-n[np.linalg.norm(np.asarray(gates)[:,None]-cent[None],axis=2).argmin(axis=1)]
  directions/=np.linalg.norm(directions,axis=1,keepdims=True)
 # Direction means incoming metal flow. The inlet lies upstream of the attachment.
 inlets=np.asarray(gates)-directions*max(12.,radius*3)
 spec={'geometry':{'attachment_points_mm':np.asarray(gates).tolist(),'inlet_centers_mm':inlets.tolist(),'inlet_directions':directions.tolist(),'gate_radius_mm':radius,'gate_count':count},'process':{'inlet_velocity_m_s':.5}}

 tet=volume['tetrahedra'];boundary=volume['triangles'];part_nodes=len(points)
 if ck.get('architecture')=='StreamNet':
  from inlet_display_mesh import append_inlets
  points,tet,boundary=append_inlets(points,tet,boundary,spec['geometry'])
 data={'points_m':points/1000,'encoder_surface_mm':surface,'encoder_normals':a['encoder_normals']@r.T,'tetrahedra':tet,'volume_triangles':boundary,'part_nodes':part_nodes}

 return a,v,points,spec,data,{'rotation':r,'translation':t,'gravity':g}



def predict(j,opt):

 start=time.perf_counter();a,v,p,spec,data,axis=config(j,opt);cloud,q=features(spec,data);cloud=torch.from_numpy(cloud);q=torch.from_numpy(q);features_done=time.perf_counter()

 fill_end=float(opt.get('fill_end',4));cool_end=float(opt.get('cool_end',1800))

 if not (0<fill_end<=30 and 0<cool_end<=20000):raise ValueError('Time ranges exceed app limits.')

 def head(task,t,code):

  y=model(cloud,q,torch.full((len(q),),float(t),device=DEVICE),task,code).cpu().numpy()

  return y*np.array(ck['stats'][task]['std'])+np.array(ck['stats'][task]['mean'])

 with lock,torch.inference_mode():
  model_start=time.perf_counter();queue_seconds=model_start-features_done
  cloud=cloud.to(DEVICE);q=q.to(DEVICE)
  code=model.encode(cloud);maps=head('maps',0,code);mechanical=head('mechanical',0,code)

  if 'cool_end' not in opt:
   # Select a display window from learned SF, never force cells to solidify.
   part_count=data.get('part_nodes',len(q))
   cool_end=float(np.clip(np.nanquantile(maps[:part_count,1],.99)*1.25,30,20000))
   for _ in range(10):
    probe=head('cooling',cool_end,code)[:part_count,1]
    if np.mean(probe>=.95)>=.99 or cool_end>=20000:break
    cool_end=min(20000,cool_end*2)
  ft=np.linspace(0,fill_end,65 if ck.get('architecture')=='StreamNet' else 25);ct=cool_end*np.linspace(0,1,25)**2

  filling=np.stack([head('filling',t,code) for t in ft]);cooling=np.stack([head('cooling',t,code) for t in ct])

  phases=np.linspace(0,1,25);deformation=None

  if warp_model is not None:

   final_mm=torch.tensor(mechanical[:,:3]*1000,dtype=torch.float32,device=DEVICE)

   deformation=np.stack([warp_model(q,code,torch.full((len(q),),float(p),device=DEVICE),final_mm).cpu().numpy()/1000 for p in phases])

 model_done=time.perf_counter()
 # All vector fields are rotated back together with the display mesh to the uploaded model frame.

 r=axis['rotation'];t=axis['translation'];gravity=axis['gravity']

 mechanical[:,:3]=mechanical[:,:3]@r;filling[:,:,2:5]=filling[:,:,2:5]@r

 if deformation is not None:deformation=deformation@r

 v=a['vertices'];p=(data['points_m']*1000-t)@r

 gate_model=(np.asarray(spec['geometry']['attachment_points_mm'])-t)@r

 def clean(x):
  a=np.asarray(x).round(6)
  return a if opt.get('_binary') and a.ndim>=2 and a.size>100 else a.tolist()

 result={'id':j,'name':jobs[j]['name'],'geometry':jobs[j]['geometry'],'known_source':jobs[j]['known_source'],'axis':int(np.argmax(np.abs(gravity))),'spec':spec,'model':model_info(),'coordinate_space':'model','gravity_direction':clean(gravity),'gate_points_model_mm':clean(gate_model),'casting_transform':{'rotation':clean(r),'translation_mm':clean(t)},'gate_extrapolation':spec['geometry']['gate_count']>2,

 'vertices':clean(v),'triangles':a['triangles'].tolist(),'points':clean(p),'volume_triangles':data['volume_triangles'].tolist(),'tetrahedra':data['tetrahedra'].tolist(),'display_mesh':json.loads((STORE/j/'volume.json').read_text()),'maps':clean(maps),'mechanical':clean(mechanical),

 'filling':clean(filling),'cooling':clean(cooling),'fill_times':clean(ft),'cool_times':clean(ct),'inference_and_features_ms':(time.perf_counter()-start)*1000,

 'timings':{'geometry_features_s':features_done-start,'queue_s':queue_seconds,'model_s':model_done-model_start},
 'cooling_window_source':'User-specified' if 'cool_end' in opt else 'Predicted solidification map and transient SF; 99% of part nodes at SF >= 0.95, capped at 20000 s',
 'checkpoint_sha256':hashlib.sha256(CK.read_bytes()).hexdigest(),'model_version':CK.parent.name,'source_sha256':jobs[j]['sha256'],

 'limitations':['Unvalidated model; unseen geometry performance unknown','Gates are feature locations, not Boolean-validated runner solids','Fixed 17-4PH,1600C pour,800C shell,6mm shell,0.5m/s inlet','Final displacement at hot demolding; no predicted warp timeline or room-temperature warp','Mesh rendering interpolates model outputs; denser display does not establish accuracy','Stress tensor head omitted until native tensor export is verified','Missing physical constraints: predictions can be nonphysical']}

 if deformation is not None:

  result.update(deformation=clean(deformation),deformation_phases=clean(phases),deformation_checkpoint_sha256=hashlib.sha256(WARP.read_bytes()).hexdigest(),deformation_scope='Learned normalized thermomechanical + demolding progress; absolute duration is not predicted. Start/end constraints are imposed.')

  result['limitations'][3]='Deformation trajectory uses normalized cooling/demolding progress; no predicted absolute duration or room-temperature warp'

 if ck.get('architecture')=='StreamNet':
  result.update(liquid_distance=clean(filling[:,:,-1]),liquid_distance_units='m',liquid_distance_sign='negative liquid, positive air',liquid_surface_method='Learned transient signed distance; zero isosurface interpolated in display tetrahedra',inlet_direction_model=clean(np.asarray(spec['geometry']['inlet_directions'])@r))
  result['limitations'][1]='Straight inlet stubs are display/query geometry; manufactured gate joints are not CAD-validated at inference.'
  result['part_node_count']=data['part_nodes']
  result['velocity_display_max']=float(max(.5,np.quantile(np.linalg.norm(filling[:,:,2:5],axis=-1),.995)))

 jobs[j]['last_setup']={'gates':result['gate_points_model_mm'],'gravity':result['gravity_direction'],'gate_directions':result.get('inlet_direction_model'),'radius':spec['geometry']['gate_radius_mm']}
 (STORE/j/'job.json').write_text(json.dumps(jobs[j],allow_nan=False))
 if not opt.get('_binary'):(STORE/j/'prediction.json').write_text(json.dumps(result,allow_nan=False))
 return result

def binary_prediction(result):
 # Header followed by aligned little-endian arrays; no millions of JSON floats.
 import struct
 meta=dict(result);chunks=[];arrays={};offset=0
 for key in ['vertices','triangles','points','volume_triangles','tetrahedra','maps','mechanical','filling','cooling','deformation','liquid_distance']:
  if key not in meta:continue
  a=np.asarray(meta.pop(key),dtype='<i4' if key in ('triangles','volume_triangles','tetrahedra') else '<f4')
  if not np.isfinite(a).all():raise ValueError('Nonfinite prediction field: '+key)
  arrays[key]={'shape':a.shape,'dtype':'i4' if a.dtype.kind=='i' else 'f4','offset':offset,'length':a.size}
  chunk=a.tobytes();chunks.append(chunk);offset+=len(chunk)
 header=json.dumps({'metadata':meta,'arrays':arrays},allow_nan=False,separators=(',',':')).encode()
 header+=b' '*((-len(header))%4)
 return struct.pack('<I',len(header))+header+b''.join(chunks)

class Handler(SimpleHTTPRequestHandler):

 def __init__(self,*args,**kw):super().__init__(*args,directory=str(ROOT/'artifacts/casting_app'),**kw)

 def end_headers(self):
  self.send_header('Cache-Control','no-store')
  super().end_headers()

 def reply(self,obj,code=200):

  b=json.dumps(obj,allow_nan=False).encode();compressed=False
  if len(b)>65536 and 'gzip' in self.headers.get('Accept-Encoding',''):
   import gzip
   b=gzip.compress(b,compresslevel=1);compressed=True
  self.send_response(code);self.send_header('Content-Type','application/json');self.send_header('Content-Length',str(len(b)))
  if compressed:self.send_header('Content-Encoding','gzip');self.send_header('Vary','Accept-Encoding')
  self.end_headers();self.wfile.write(b)

 def do_GET(self):
  if self.path=='/evaluation':
    b=CK.with_name('report.json').read_bytes();self.send_response(200);self.send_header('Content-Type','application/json');self.send_header('Content-Length',str(len(b)));self.end_headers();self.wfile.write(b);return

  if self.path=='/api/model':return self.reply(model_info())

  if self.path.startswith('/api/geometry/'):

   j=self.path.rsplit('/',1)[-1]

   if j not in jobs or jobs[j]['status']!='ready':return self.reply({'error':'Geometry is not ready'},404)

   a=np.load(STORE/j/'geometry.npz')

   setup=jobs[j].get('last_setup')
   if setup and not setup.get('gate_directions'):
    import trimesh
    tri=a['vertices'][a['triangles']];normals=np.cross(tri[:,1]-tri[:,0],tri[:,2]-tri[:,0]);valid=np.linalg.norm(normals,axis=1)>1e-10;tri=tri[valid];normals=normals[valid];directions=[]
    for point in setup['gates']:
     closest=trimesh.triangles.closest_point(tri,np.broadcast_to(point,(len(tri),3)));idx=np.argmin(np.linalg.norm(closest-point,axis=1));directions.append((-normals[idx]/np.linalg.norm(normals[idx])).tolist())
    setup['gate_directions']=directions
   return self.reply({**jobs[j],'vertices':a['vertices'].tolist(),'triangles':a['triangles'].tolist(),'coordinate_space':'model'})

  if self.path=='/api/recent':

   ordered=sorted(jobs.values(),key=lambda j:(STORE/j['id']/'input.step').stat().st_mtime,reverse=True)

   return self.reply([{'id':j['id'],'name':j['name']} for j in ordered if 'bracket' not in j['name'].lower()][:6])

  if self.path.startswith('/api/job/'):

   j=self.path.rsplit('/',1)[-1];return self.reply(jobs.get(j,{'error':'Unknown job'}),200 if j in jobs else 404)

  return super().do_GET()

 def do_POST(self):

  try:

   size=int(self.headers.get('Content-Length',0))

   if size<=0 or size>50*1024*1024:raise ValueError('Upload limit is 50 MB.')

   raw=self.rfile.read(size)

   if self.path=='/api/resume':

    j=json.loads(raw)['id']

    if j not in jobs:raise ValueError('Unknown saved part')

    if jobs[j]['status']!='ready':jobs[j]['status']='processing';threading.Thread(target=prepare,args=(j,),daemon=True).start()

    return self.reply(jobs[j])

   if self.path=='/api/upload':

    if b'ISO-10303-21' not in raw[:4096].upper():raise ValueError('Expected a STEP/STP file.')

    j=uuid.uuid4().hex;folder=STORE/j;folder.mkdir();(folder/'input.step').write_bytes(raw);sha=hashlib.sha256(raw).hexdigest()

    jobs[j]={'id':j,'status':'processing','name':self.headers.get('X-Filename','Uploaded STEP'),'sha256':sha,'known_source':KNOWN.get(sha)}

    threading.Thread(target=prepare,args=(j,),daemon=True).start();return self.reply(jobs[j])

   if self.path=='/api/predict':

    opt=json.loads(raw);j=opt['id']

    if j not in jobs or jobs[j]['status']!='ready':raise ValueError('Upload geometry is not ready.')

    if 'application/x-ace-arrays' in self.headers.get('Accept',''):
     opt['_binary']=True;result=predict(j,opt);b=binary_prediction(result)
     (STORE/j/'prediction.acebin').write_bytes(b)
     self.send_response(200);self.send_header('Content-Type','application/x-ace-arrays');self.send_header('Content-Length',str(len(b)));self.end_headers();self.wfile.write(b);return
    return self.reply(predict(j,opt))

   self.reply({'error':'Not found'},404)

  except Exception as e:self.reply({'error':str(e)},400)

if __name__=='__main__':

 if DEVICE.type=='cuda':
  with torch.inference_mode():
   cloud=torch.zeros((256,6),device=DEVICE);q=torch.zeros((256,ck['query_dim']),device=DEVICE);t=torch.zeros(256,device=DEVICE);code=model.encode(cloud)
   for task in ck['stats']:model(cloud,q,t,task,code)
   if warp_model is not None:warp_model(q,code,t,torch.zeros((256,3),device=DEVICE))
   torch.cuda.synchronize()
 port=int(os.environ.get('ACE_PORT','8767'));print(f'Ace-Mini http://127.0.0.1:{port}',flush=True);ThreadingHTTPServer(('127.0.0.1',port),Handler).serve_forever()