"""Prepared graph_6000 (5,925 actual nodes) weak-label diagnostic: normal baseline vs dihedral GNN."""
from __future__ import annotations
import argparse,json,sys
from pathlib import Path
import numpy as np
from scipy.spatial import cKDTree
import torch
from torch import nn
import torch.nn.functional as F
import matplotlib;matplotlib.use('Agg')
import matplotlib.pyplot as plt
from models.data import ROOT,fixed_physical_normalize,split_embedded_units
from models.training import seed_everything
EXP=ROOT/'dihedral_gnn';INPUT=EXP/'input_processed';MODEL=EXP/'model';OUT=EXP/'outputs';SCALE=30.
NAMES=np.asarray(['pos_norm_x','pos_norm_y','pos_norm_z','normal_x','normal_y','normal_z','vertex_area_mm2','mean_incident_angle_rad','mean_curvature_abs_1_per_mm','gaussian_curvature_abs_1_per_mm2','boundary_node','feature_node','topological_degree','analytic_k1_abs_1_per_mm','analytic_k2_abs_1_per_mm','analytic_gaussian_signed_1_per_mm2','cad_k1_signed_1_per_mm','cad_k2_signed_1_per_mm','cad_mean_H_signed_1_per_mm','cad_gaussian_K_1_per_mm2','cad_curvedness_1_per_mm','cad_shape_index'])
ENAMES=np.asarray(['length_mm','direction_x','direction_y','direction_z','normal_dot','dihedral_rad','boundary_edge','feature_edge'])

def cg():
 p=str(ROOT/'전처리');sys.path.insert(0,p) if p not in sys.path else None
 import cad_graphs
 return cad_graphs
def graph():
 """One prepared ~6k CAD graph only: requested 6,000, actual 5,925 nodes."""
 source=ROOT/'data_processed/graphs/graph_6000.npz'
 with np.load(source,allow_pickle=False) as a:z={k:a[k] for k in a.files}
 v=z['pos_mm'].astype('f4');n=z['normal'].astype('f4');x=z['x'].astype('f4');mask=z['feature_valid'].astype(bool);edge=z['edge_index'].astype('i8');eraw=z['edge_attr'].astype('f4');ea=eraw.copy();ea[:,0]/=1000;ea[:,5]/=np.pi
 if len(v)!=5925:raise RuntimeError(f'graph_6000 expected 5925 actual nodes, got {len(v)}')
 c=np.load(ROOT/'analysis/results/cad_surface_curvature.npz',allow_pickle=False);d,ix=cKDTree(c['points_mm']).query(v,workers=1);dot=(n*c['parameter_normals'][ix]).sum(1);ok=(d<3)&(np.abs(dot)>.8)&np.isfinite(c['k1_per_mm'][ix])&np.isfinite(c['k2_per_mm'][ix]);orient=np.where(dot>=0,1.,-1.);a=c['k1_per_mm'][ix]*orient;b=c['k2_per_mm'][ix]*orient;k1=np.maximum(a,b);k2=np.minimum(a,b);H=(k1+k2)/2;K=k1*k2;C=np.sqrt((k1*k1+k2*k2)/2);S=(2/np.pi)*np.arctan2(k1+k2,k1-k2);extra=np.c_[k1,k2,H,K,C,S].astype('f4');extra[~ok]=0;x=np.c_[x,extra];mask=np.c_[mask,np.repeat(ok[:,None],6,1)]
 INPUT.mkdir(parents=True,exist_ok=True);aug=dict(z);aug.update(x=x.astype('f4'),feature_valid=mask,feature_names=NAMES.astype('<U48'));np.savez_compressed(INPUT/'prepared_graph_5925.npz',**aug);np.savez_compressed(INPUT/'dihedral_signed_curvature_features.npz',cad_k1_signed_1_per_mm=k1.astype('f4'),cad_k2_signed_1_per_mm=k2.astype('f4'),cad_mean_H_signed_1_per_mm=H.astype('f4'),cad_gaussian_K_1_per_mm2=K.astype('f4'),cad_curvedness_1_per_mm=C.astype('f4'),cad_shape_index=S.astype('f4'),orientation_valid=ok,mapping_distance_mm=d.astype('f4'),normal_dot=dot.astype('f4'))
 meta={'requested_node_budget':6000,'actual_node_count':5925,'triangle_count':int(z['face'].shape[1]),'source_graph':'data_processed/graphs/graph_6000.npz','topology':'saved edge_index derives from saved faces; face dihedral remains the two-adjacent-face quantity','node_padding_or_duplication':False,'signed_curvature_mapping':{'coverage_fraction':float(ok.mean()),'valid_count':int(ok.sum()),'invalid_count':int((~ok).sum()),'valid_rule':'nearest STEP analytic sample <3 mm and abs(prepared normal dot STEP parameter normal) >0.8; both k values are multiplied by sign(dot) into prepared-normal convention, then ordered k1>=k2; invalid signs are masked','formulas':'H=(k1+k2)/2; K=k1*k2; curvedness=sqrt((k1²+k2²)/2); shape-index=(2/pi)*atan2(k1+k2,k1-k2)','caveat':'sign is meaningful only under this local oriented correspondence; use dihedral at sharp creases, where smooth curvature is undefined'}};(INPUT/'prepared_graph_5925_provenance.json').write_text(json.dumps(meta,ensure_ascii=False,indent=2),encoding='utf8')
 return {'pos':v,'normal':n,'raw':x,'node':np.c_[fixed_physical_normalize(x,NAMES),mask.astype('f4')].astype('f4'),'edge':edge,'eraw':eraw,'ea':ea,'face':z['face'].astype('i4'),'meta':meta}
def pairs():
 out=[]
 for p in sorted((ROOT/'data_processed/paired').glob('*_cad_ply_targets.npz')):
  with np.load(p,allow_pickle=False) as a:out.append((str(a['run_id'].item()),{k:a[k] for k in a.files}))
 if len(out)!=2:raise RuntimeError('requires exactly two paired runs')
 return out
def proc(x):
 raw=x['process_x'].astype('f4').copy();obs=x['process_feature_valid'].astype(bool);names,units=split_embedded_units(x['process_feature_names'].astype(str));names=names.tolist();imp=np.zeros(len(raw),bool)
 for pre,val,flag in [('상부온도',x['heater_top_c'],x['heater_top_imputed']),('하부온도',x['heater_bottom_c'],x['heater_bottom_imputed'])]:
  for z in range(1,18):
   k=f'{pre}{z:02d}'
   if k in names and flag[z-1]:raw[names.index(k)]=val[z-1];imp[names.index(k)]=1
 good=obs|imp;return np.r_[fixed_physical_normalize(np.where(good,raw,0)[None],np.asarray(names),units)[0],obs.astype('f4'),imp.astype('f4')],names
def labels(g,rows):
 r=[];d=[];dist=[]
 for _,x in rows:
  q,i=cKDTree(x['cad_pos_mm']).query(g['pos'],workers=1);dist.append(q.astype('f4'));r.append(x['weak_defect_probability'][i,1].astype('f4'));d.append(np.clip(x['cad_to_scan_normal_residual_mm'][i],-SCALE,SCALE).astype('f4')/SCALE)
 return r,d,dist
class Net(nn.Module):
 def __init__(self,a,b,c,w=48):
  super().__init__();self.node=nn.Sequential(nn.Linear(a,w),nn.GELU(),nn.Linear(w,w),nn.GELU());self.edge=nn.Sequential(nn.Linear(b,w),nn.GELU(),nn.Linear(w,w));self.process=nn.Sequential(nn.Linear(c,w),nn.GELU(),nn.Linear(w,w));self.msg=nn.Sequential(nn.Linear(4*w,w),nn.GELU(),nn.Linear(w,w),nn.GELU());self.up=nn.ModuleList([nn.GRUCell(2*w,w) for _ in range(2)]);self.risk=nn.Linear(w,1);self.signed=nn.Linear(w,1);self.mag=nn.Linear(w,1)
 def forward(self,x,e,ea,p,b):
  h=self.node(x);q=self.process(p)[b];s,t=e;z=self.edge(ea)
  for u in self.up:m=self.msg(torch.cat((h[s],h[t]-h[s],z,q[s]),1));a=torch.zeros_like(h).index_add_(0,t,m);h=u(torch.cat((h,a),1),h)
  d=self.signed(h).squeeze(1);return {'risk':self.risk(h).squeeze(1),'signed':d,'mag':d.abs()+F.softplus(self.mag(h).squeeze(1))}
def metric(p,y):
 z=p>=.5;t=y>=.5;i=int((z&t).sum());u=int((z|t).sum());o=np.argsort(-p,kind='stable');a=t[o].astype(float);ap=float((np.cumsum(a)/np.arange(1,len(a)+1)*a).sum()/max(1,a.sum()));return {'AP':ap,'IoU_at_0.5':i/max(1,u),'Dice_at_0.5':2*i/max(1,int(z.sum()+t.sum())),'target_nodes':int(t.sum()),'predicted_nodes':int(z.sum())}
def select(g,kind):
 # This baseline has the same architecture, seed, epochs, loss, conditions and threshold; it excludes all explicit curvature/ridge/boundary/face-dihedral channels.
 ni=list(range(len(NAMES))) if kind=='dihedral_aware' else [0,1,2,3,4,5,6,7,12];ei=list(range(8)) if kind=='dihedral_aware' else [0,1,2,3,4]
 return np.c_[g['node'][:,ni],g['node'][:,np.asarray(ni)+len(NAMES)]],g['ea'][:,ei],[NAMES[i] for i in ni],[ENAMES[i] for i in ei]
def train(g,rows,kind,epochs,lr):
 seed_everything();x0,e0,nnames,enames=select(g,kind);risk,deform,dist=labels(g,rows);pv,_=zip(*(proc(x) for _,x in rows));n=len(g['pos']);x=np.r_[x0,x0];e=np.c_[g['edge'],g['edge']+n];ea=np.r_[e0,e0];b=np.repeat(np.arange(2),n);tr=np.r_[risk[0],risk[1]];td=np.r_[deform[0],deform[1]];model=Net(x.shape[1],ea.shape[1],len(pv[0]));opt=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=1e-5);X,E,A,P,B=map(torch.as_tensor,(x,e,ea,np.asarray(pv,'f4'),b));E=E.long();B=B.long();R=torch.as_tensor(tr);D=torch.as_tensor(td);pw=((len(R)-R.sum())/R.sum().clamp_min(1)).clamp(max=200);hist=[];best=(1e9,0,None)
 for ep in range(epochs):
  opt.zero_grad();o=model(X,E,A,P,B);bce=F.binary_cross_entropy_with_logits(o['risk'],R,pos_weight=pw);dice=1-(2*(torch.sigmoid(o['risk'])*R).sum()+1)/(torch.sigmoid(o['risk']).sum()+R.sum()+1);dl=F.huber_loss(o['signed'],D,delta=.15);ml=F.huber_loss(o['mag'],D.abs(),delta=.15);loss=bce+dice+.2*dl+.2*ml;loss.backward();opt.step();h={'epoch':ep+1,'total':float(loss.detach()),'risk_bce':float(bce.detach()),'risk_dice':float(dice.detach()),'signed_deformation_huber':float(dl.detach()),'magnitude_deformation_huber':float(ml.detach())};hist.append(h)
  if h['total']<best[0]:best=(h['total'],ep+1,{k:v.detach().clone() for k,v in model.state_dict().items()})
 model.load_state_dict(best[2]);model.eval()
 with torch.no_grad():o=model(X,E,A,P,B);p=torch.sigmoid(o['risk']).numpy();s=o['signed'].numpy()*SCALE;m=o['mag'].numpy()*SCALE
 runs=[{'run':run,'risk':p[j*n:(j+1)*n],'signed':s[j*n:(j+1)*n],'mag':m[j*n:(j+1)*n],'weak':risk[j],'target_signed':deform[j]*SCALE,'map_distance_mm':dist[j],'metric':metric(p[j*n:(j+1)*n],risk[j])} for j,(run,_) in enumerate(rows)]
 return model,{'variant':kind,'node_features':list(nnames),'edge_features':list(enames),'history':hist,'best_loss':best[0],'best_epoch':best[1],'runs':runs}
def view(g,infos,rows):
 n=len(g['pos']);pick=np.linspace(0,g['edge'].shape[1]-1,min(4500,g['edge'].shape[1]),dtype=int);total=np.zeros(n);degree=np.zeros(n);np.add.at(total,g['edge'][0],g['eraw'][:,5]);np.add.at(degree,g['edge'][0],1);first=np.full(n,-1,'i4')
 for j,a in enumerate(g['edge'][0]):
  if first[a]<0:first[a]=j
 inc=g['edge'][:,first].T.tolist();faces=g['face'].T;fm={}
 for fi,t in enumerate(faces):
  for a,b in ((t[0],t[1]),(t[1],t[2]),(t[2],t[0])):fm.setdefault(tuple(sorted((int(a),int(b)))),[]).append(fi)
 tri=g['pos'][faces];fn=np.cross(tri[:,1]-tri[:,0],tri[:,2]-tri[:,0]);fn/=np.maximum(np.linalg.norm(fn,axis=1,keepdims=True),1e-8);fg=[]
 for a,b in inc:
  fs=fm.get(tuple(sorted((a,b))),[])
  if len(fs)==2:fg.append({'status':'two_face','triangles':tri[fs].round(4).tolist(),'normals':fn[fs].round(6).tolist(),'theta_rad':float(np.arccos(np.clip(fn[fs[0]]@fn[fs[1]],-1,1)))})
  else:fg.append({'status':'boundary' if len(fs)==1 else 'nonmanifold_or_missing','face_count':len(fs),'policy':'face dihedral undefined; no theta invented'})
 models=[]
 for info in infos:
  rs=[]
  for r,(_,x) in zip(info['runs'],rows):
   names=split_embedded_units(x['process_feature_names'].astype(str))[0].tolist();keys=[]
   for k in ('가열시간','가열온도','상부원단온도','하부원단온도'):
    j=names.index(k);keys.append({'name':k,'value':float(x['process_x'][j]) if x['process_feature_valid'][j] else None,'observed':bool(x['process_feature_valid'][j])})
   rs.append({**{k:(v.round(6).tolist() if isinstance(v,np.ndarray) else v) for k,v in r.items() if k in ('run','risk','signed','mag','weak','metric')},'magnitude':r['mag'].round(5).tolist(),'heater_top_c':x['heater_top_c'].round(2).tolist(),'heater_bottom_c':x['heater_bottom_c'].round(2).tolist(),'heater_top_imputed':x['heater_top_imputed'].astype(bool).tolist(),'heater_bottom_imputed':x['heater_bottom_imputed'].astype(bool).tolist(),'key_global_process':keys})
  models.append({'id':info['variant'],'label':'Dihedral-aware GNN' if info['variant']=='dihedral_aware' else 'Normal GNN baseline','best_loss':info['best_loss'],'runs':rs})
 raw=g['raw'];valid=g['node'][:,len(NAMES)+18].astype(bool);shapevalid=valid&(raw[:,20]>1e-6);shape_class=np.where(~valid,0,np.where(raw[:,19]<0,1,np.where(~shapevalid,2,np.where(raw[:,18]>=0,3,4))));return {'nodes':g['pos'].round(4).tolist(),'normals':g['normal'].round(5).tolist(),'edges':g['edge'][:,pick].T.tolist(),'node_feature':{'mean_curvature_abs_1_per_mm':raw[:,8].round(6).tolist(),'feature_node':raw[:,11].round(0).tolist(),'boundary_node':raw[:,10].round(0).tolist(),'cad_mean_H_signed_1_per_mm':raw[:,18].round(6).tolist(),'cad_shape_index':raw[:,21].round(6).tolist(),'cad_signed_curvature_valid':valid.tolist(),'cad_shape_index_valid':shapevalid.tolist(),'cad_shape_class':shape_class.tolist(),'mean_adjacent_dihedral_rad':(total/np.maximum(degree,1)).round(6).tolist()},'incident_edges':inc,'incident_edge_feature':{str(k):g['eraw'][first,i].round(6).tolist() for i,k in enumerate(ENAMES)},'incident_face_geometry':fg,'models':models,'edge_draw_note':f'4,500 of {g["edge"].shape[1]} directed edges drawn; selected edge and all-edge mean are exact.','limitations':['Prepared graph_6000: requested ~6k, actual 5,925 nodes; no padding or duplication.','Signed CAD curvature is gray/unknown outside the orientation-valid mask; shape index is additionally undefined at near-zero curvedness. H positive/negative follows prepared-normal convention, not universal outward convexity; K<0 is saddle.','Two NG scans and point-dilated weak labels only; no held-out generalization.','No validated registration, local thermal field, material/contact/FLC, stress/strain, or tear physics.']}
def run(a):
 if not a.allow_diagnostic_weak_labels:raise SystemExit('Pass --allow-diagnostic-weak-labels: targets are weak/unvalidated.')
 seed_everything();g=graph();rows=pairs();infos=[]
 for k in ('normal_baseline','dihedral_aware'):
  m,i=train(g,rows,k,a.epochs,a.learning_rate);MODEL.mkdir(parents=True,exist_ok=True);torch.save({'state_dict':m.state_dict(),'variant':k,'node_count':len(g['pos']),'node_features':i['node_features'],'edge_features':i['edge_features'],'production_valid':False},MODEL/f'{k}_gnn.pt');infos.append(i)
 OUT.mkdir(parents=True,exist_ok=True);d={x['variant']:x for x in infos};np.savez_compressed(OUT/'prepared_graph_5925_direct_predictions.npz',cad_pos_mm=g['pos'],cad_normal=g['normal'],node_count=np.int32(len(g['pos'])),run_id=np.asarray([x['run'] for x in infos[0]['runs']]),normal_baseline_risk_probability=np.stack([x['risk'] for x in d['normal_baseline']['runs']]),dihedral_risk_probability=np.stack([x['risk'] for x in d['dihedral_aware']['runs']]),normal_baseline_signed_mm=np.stack([x['signed'] for x in d['normal_baseline']['runs']]),dihedral_signed_mm=np.stack([x['signed'] for x in d['dihedral_aware']['runs']]),normal_baseline_magnitude_mm=np.stack([x['mag'] for x in d['normal_baseline']['runs']]),dihedral_magnitude_mm=np.stack([x['mag'] for x in d['dihedral_aware']['runs']]),direct_model_evaluation=np.bool_(True))
 for p in (OUT/'direct_original_cad_predictions.npz',OUT/'fixed_6000_direct_predictions.npz',OUT/'direct_original_predictions.png',OUT/'training.png',INPUT/'finest_prepared_cad_graph_11956.npz',INPUT/'fixed_6000_surface_graph.npz',INPUT/'fixed_6000_surface_provenance.json'):
  if p.exists():p.unlink()
 contract={'surface_graph':'data_processed/graphs/graph_6000.npz','requested_node_budget':6000,'actual_node_count':len(g['pos']),'triangle_count':int(g['face'].shape[1]),'training_inference_viewer_outputs':'same prepared graph_6000 with 5,925 actual nodes','node_density_is_risk_feature':False,'normal_baseline':'same GNN/seed/epochs/loss/process/threshold; excludes curvature, analytic curvature, boundary/ridge and face dihedral','dihedral_face_definition':'acos(clamp(dot(two adjacent face normals), -1,1)); undefined on boundary/nonmanifold','changed_cad':'A new CAD can execute with variable node count/topology, but this experiment fixes the ~6k preparation rule: independently simplify each new CAD with the same rule; node indices never correspond across designs; output is direct per-new-CAD-node. Generalization is unmeasured with one CAD/two runs. Use multi-design training, CAD-family held-out splits and OOD uncertainty before interpreting cross-design output.','features':{x['variant']:{'node':x['node_features'],'edge':x['edge_features']} for x in infos}};(INPUT/'feature_contract.json').write_text(json.dumps(contract,ensure_ascii=False,indent=2),encoding='utf8');(OUT/'viewer_data.json').write_text(json.dumps(view(g,infos,rows),ensure_ascii=False,separators=(',',':')),encoding='utf8')
 fig,ax=plt.subplots(figsize=(8,3.5));[ax.plot([h['epoch'] for h in x['history']],[h['total'] for h in x['history']],label=x['variant']) for x in infos];ax.set_yscale('log');ax.legend();ax.grid(alpha=.2);ax.set_title('Same prepared graph_6000 (5,925 actual nodes) weak-label fit');fig.savefig(OUT/'loss_comparison.png',dpi=160);plt.close(fig)
 fig,axes=plt.subplots(2,2,figsize=(10,8));
 for i,x in enumerate(infos):
  for j,r in enumerate(x['runs']):c=axes[i,j].scatter(g['pos'][:,0],g['pos'][:,2],c=r['risk'],s=2,cmap='magma',vmin=0,vmax=1);axes[i,j].set_title(f"{x['variant']} {r['run']} risk");axes[i,j].set_aspect('equal');fig.colorbar(c,ax=axes[i,j])
 fig.savefig(OUT/'fixed_6000_model_overlays.png',dpi=160);plt.close(fig)
 comp=[{'run':a['run'],'normal_baseline':a['metric'],'dihedral_aware':b['metric'],'dihedral_better_AP':b['metric']['AP']>a['metric']['AP'],'dihedral_better_IoU':b['metric']['IoU_at_0.5']>a['metric']['IoU_at_0.5']} for a,b in zip(d['normal_baseline']['runs'],d['dihedral_aware']['runs'])];concept={'status':'conceptual_next_model_not_trained','branches':'CAD graph geometry encoder (curvature/face-dihedral/ridge/boundary) + XLS thermal/process encoder; local thermal projection only when heater footprints exist','node_heads':'segmentation/risk/deformation/thinning','region_heads':'defect type and physical area mm²','identifiability':'thermal risk and geometry risk cannot be separated from one fused label without branch-specific targets, controlled variation or concept supervision','current_data':'one CAD, two NG runs, only 찍힘, no observed area; weak dilation is not physical area','physics_needed':'local temperature, thickness/material/FLC/contact and measured strain/stress or simulation targets are required for causal sharp-bump→strain→tear claims'};report={'status':'completed','production_valid':False,'fixed_graph':'same prepared graph_6000 representation (requested ~6k; 5,925 actual nodes) for training, inference, viewer and outputs','comparison':comp,'models':[{'variant':x['variant'],'best_loss':x['best_loss'],'best_epoch':x['best_epoch'],'loss_terms':'BCE + Dice + 0.2×signed-normal Huber + 0.2×magnitude Huber against abs(normal residual)','history':x['history'],'metrics':[r['metric'] for r in x['runs']]} for x in infos],'inputs':contract,'changed_cad':contract['changed_cad'],'conceptual_next_model':concept,'limitations':['diagnostic weak-label fit; not a physical tear/defect predictor']};(OUT/'report.json').write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding='utf8');return report
def main():
 p=argparse.ArgumentParser();p.add_argument('--allow-diagnostic-weak-labels',action='store_true');p.add_argument('--epochs',type=int,default=56);p.add_argument('--learning-rate',type=float,default=2e-3);a=p.parse_args();print(json.dumps(run(a),ensure_ascii=False,indent=2))
if __name__=='__main__':main()