"""Run actual-schema forward/backward smokes; real training stays label-gated."""
from __future__ import annotations
import argparse, json, os, time
from pathlib import Path
from typing import Any, Callable
import numpy as np
import torch
from torch import Tensor, nn
from .architectures import CNN2DUNet, UNet3D, SparseUNetOptional, PointNetSeg, DGCNNSeg, MeshGraphProcessor, PhysicsConstrainedRiskGNN, UVNetFaceModel, FNO2d, TabularSurrogate
from .data import ROOT, InputContractError, image_input, voxel_input, sparse_voxel_input, sparse_batch_collate, point_inputs, graph_inputs, brep_input, tabular_input, process_conditions
from .training import seed_everything, fit_validated_npz, validated_contract_dimensions, DatasetContractError

OUT=ROOT/"results/models"; torch.set_num_threads(max(1,min(4,os.cpu_count() or 1)))

def _tensor(a: np.ndarray, dtype=torch.float32) -> Tensor: return torch.as_tensor(a,dtype=dtype)
def _output_tensors(y: Tensor | dict[str,Tensor]) -> list[Tensor]:
    values=[y] if isinstance(y,Tensor) else list(y.values())
    if not values or any(not isinstance(value,Tensor) for value in values):
        raise RuntimeError("model output must be a tensor or a non-empty tensor dictionary")
    return values

def _finite(y: Tensor | dict[str,Tensor]) -> None:
    if any(not torch.isfinite(value).all() for value in _output_tensors(y)): raise RuntimeError("model output is non-finite")

def _smoke(family: str, model: nn.Module, forward: Callable[[],Tensor|dict[str,Tensor]], source: str, role: str, extra: dict[str,Any]|None=None) -> dict[str,Any]:
    """One real-input forward + synthetic-gradient backward. No targets are implied."""
    model.train(); model.zero_grad(set_to_none=True); t=time.perf_counter(); y=forward(); fwd=(time.perf_counter()-t)*1000; _finite(y)
    t=time.perf_counter(); synthetic_objective=sum(value.square().mean() for value in _output_tensors(y)); synthetic_objective.backward(); bwd=(time.perf_counter()-t)*1000
    grads=[p.grad for p in model.parameters() if p.requires_grad]
    if not grads or any(g is None or not torch.isfinite(g).all() for g in grads): raise RuntimeError("non-finite/missing parameter gradient")
    shape={key:list(value.shape) for key,value in y.items()} if isinstance(y,dict) else list(y.shape)
    record={"family":family,"status":"forward_backward_smoke","source":source,"production_role":role,"params":sum(p.numel() for p in model.parameters()),"forward_latency_ms_cpu":round(fwd,3),"backward_latency_ms_cpu":round(bwd,3),"output_shape":shape,"finite_output":True,"finite_parameter_gradients":True,"synthetic_backward_objective":"sum of mean(square(untrained output tensors)); not a label, fit, or accuracy","accuracy":None}
    if extra: record.update(extra)
    return record

def _synthetic_checkpoint(name: str, model: nn.Module, forward: Callable[[],Tensor|dict[str,Tensor]]) -> dict[str,Any]:
    """Two-step synthetic-only trainability check, kept separate from actual-data status."""
    model.train(); opt=torch.optim.Adam(model.parameters(),lr=1e-4); losses=[]
    for _ in range(2):
        opt.zero_grad(); out=forward()
        loss=sum(value.square().mean() for value in _output_tensors(out))  # explicitly synthetic, not AC28 labels
        loss.backward(); opt.step(); losses.append(float(loss.detach()))
    path=OUT/"checkpoints"/f"synthetic_{name}.pt"; path.parent.mkdir(parents=True,exist_ok=True)
    torch.save({"synthetic_only":True,"target":"zeros_like_model_output","steps":2,"state_dict":model.state_dict()},path)
    return {"family":name,"status":"synthetic_learnability_smoke","synthetic_only":True,"steps":2,"losses":losses,"loss_decreased_or_equal":bool(losses[-1] <= losses[0]),"checkpoint":str(path.relative_to(ROOT)),"accuracy":None}

def _conditions() -> tuple[Tensor,list[str],str]:
    x,m,runs,source=process_conditions()
    return _tensor(np.concatenate([x,m.astype(np.float32)],1)),runs,source

def _image_families(synth: list[dict[str,Any]]) -> list[dict[str,Any]]:
    d=image_input(); cond,runs,psource=_conditions(); base=np.concatenate([d["x"],d["mask"].astype(np.float32)],1)
    # All six actual views are evaluated under both actual process rows.
    x=_tensor(np.concatenate([base for _ in runs],0)); process=cond.repeat_interleave(base.shape[0],0)
    rec=[]; m=CNN2DUNet(x.shape[1],global_dim=process.shape[1]); rec.append(_smoke("2D_CNN_UNet",m,lambda:m(x,process),d["source"],"CAD multiview + planned process pre-process candidate only after validated labels",{"views":6,"process_runs":runs,"process_source":psource,"feature_masks_appended":True})); synth.append(_synthetic_checkpoint("2d_cnn_unet",m,lambda:m(x,process)))
    m=FNO2d(x.shape[1],global_dim=process.shape[1]); rec.append(_smoke("FNO2d",m,lambda:m(x,process),d["source"],"architecture smoke only: a real thermal/operator surrogate is unavailable until paired field inputs/targets exist",{"views":6,"process_runs":runs,"feature_masks_appended":True,"real_field_training":"unavailable"})); synth.append(_synthetic_checkpoint("fno2d",m,lambda:m(x,process)))
    return rec

def _voxel_families(synth: list[dict[str,Any]]) -> list[dict[str,Any]]:
    d=voxel_input(); cond,runs,psource=_conditions(); base=np.concatenate([d["x"],d["mask"].astype(np.float32)],0); x=_tensor(np.repeat(base[None],len(runs),0))
    m=UNet3D(x.shape[1],global_dim=cond.shape[1]); rec=[_smoke("3D_UNet",m,lambda:m(x,cond),d["source"],"CAD surface voxels + planned process pre-process candidate; sampled surface occupancy is not a filled solid",{"process_runs":runs,"process_source":psource,"feature_masks_appended":True})]; synth.append(_synthetic_checkpoint("3d_unet",m,lambda:m(x,cond)))
    try:
        import MinkowskiEngine as ME  # type: ignore
    except ImportError as exc:
        rec.append({"family":"Sparse_UNet","status":"unavailable_optional_backend","backend":"MinkowskiEngine","reason":f"backend import unavailable: {exc}","accuracy":None})
        return rec
    except Exception as exc:
        rec.append({"family":"Sparse_UNet","status":"optional_backend_import_error","backend":"MinkowskiEngine","reason":f"backend import failed unexpectedly: {type(exc).__name__}: {exc}","accuracy":None})
        return rec
    try:
        sd=sparse_voxel_input(); coordinates,features=sparse_batch_collate(sd["coords"],sd["features"])
        ct,ft=_tensor(coordinates,torch.int32),_tensor(features)
        m=SparseUNetOptional(ft.shape[1])
        def sparse_forward(m=m,ct=ct,ft=ft):
            out=m(ME.SparseTensor(features=ft,coordinates=ct))
            # A segmentation decoder must return the same sparse coordinate set;
            # this checks alignment before applying the synthetic objective to F.
            if out.C.shape != ct.shape or not torch.equal(torch.unique(out.C,dim=0),torch.unique(ct,dim=0)):
                raise RuntimeError("Sparse U-Net output coordinates do not align with input sparse voxels")
            return out.F
        rec.append(_smoke("Sparse_UNet",m,sparse_forward,sd["source"],"CAD sparse-surface geometry candidate; no process branch is implemented",{"backend":"MinkowskiEngine","sparse_voxels":int(len(ct)),"input_coordinate_shape":list(ct.shape),"output_coordinate_alignment_verified":True,"feature_masks_appended":True}))
        synth.append(_synthetic_checkpoint("sparse_unet",m,sparse_forward))
    except Exception as exc:
        rec.append({"family":"Sparse_UNet","status":"optional_backend_runtime_error","backend":"MinkowskiEngine","reason":f"backend loaded but sparse execution failed: {type(exc).__name__}: {exc}","accuracy":None})
    return rec

def _point_families(synth: list[dict[str,Any]]) -> list[dict[str,Any]]:
    rows=[]
    for d in point_inputs():
        x=_tensor(np.concatenate([d["x"],d["mask"].astype(np.float32)],1)[None])
        m=PointNetSeg(x.shape[-1]); rows.append(_smoke("PointNet_segmentation",m,lambda m=m,x=x:m(x),d["source"],"post_scan_inspection" if d["timing"]=="post_scan_diagnostic" else "CAD geometry pre-process candidate after validated outcomes",{"timing":d["timing"],"points":int(x.shape[1]),"feature_masks_appended":True})); synth.append(_synthetic_checkpoint("pointnet_"+Path(d["source"]).stem,m,lambda m=m,x=x:m(x)))
        # Explicit cap is only for DGCNN cdist memory; this is a deterministic prefix of the prepared point sample.
        z=x[:,:min(512,x.shape[1])]; m=DGCNNSeg(z.shape[-1],k=12); rows.append(_smoke("DGCNN_dynamic_feature_knn",m,lambda m=m,z=z:m(z),d["source"],"post_scan_inspection" if d["timing"]=="post_scan_diagnostic" else "CAD geometry pre-process candidate after validated outcomes",{"timing":d["timing"],"points":int(z.shape[1]),"dynamic_graphs":"feature-neighbor graph recomputed at each EdgeConv","feature_masks_appended":True})); synth.append(_synthetic_checkpoint("dgcnn_"+Path(d["source"]).stem,m,lambda m=m,z=z:m(z)))
    return rows

def _batched_graph(d: dict[str,Any], cond: Tensor) -> tuple[Tensor,Tensor,Tensor,Tensor,Tensor]:
    x=np.concatenate([d["x"],d["mask"].astype(np.float32)],1); n=len(x); b=len(cond); e=d["edge_index"]; a=np.nan_to_num(d["edge_attr"],nan=0.,posinf=0.,neginf=0.).astype(np.float32)
    # Preserve full prepared topology/attributes. First edge column is documented length_mm and gets the same fixed mm scale.
    if a.shape[1]: a[:,0]/=1000.
    node=np.concatenate([x for _ in range(b)],0); edge=np.concatenate([e+i*n for i in range(b)],1); attr=np.concatenate([a for _ in range(b)],0); batch=np.repeat(np.arange(b),n)
    return _tensor(node),_tensor(edge,torch.long),_tensor(attr),cond,_tensor(batch,torch.long)

def _procedural_geometry_probe(model: MeshGraphProcessor, cond: Tensor, node_dim: int, edge_dim: int) -> dict[str,Any]:
    """Separate procedural mesh change; topology/features are recomputed, not stale CAD fields."""
    q=8; u,v=np.meshgrid(np.linspace(-1,1,q),np.linspace(-1,1,q),indexing="ij")
    def build(amplitude: float):
        z=amplitude*np.sin(np.pi*u)*np.sin(np.pi*v); p=np.c_[u.ravel(),v.ravel(),z.ravel()].astype(np.float32)
        dzdu=amplitude*np.pi*np.cos(np.pi*u)*np.sin(np.pi*v); dzdv=amplitude*np.pi*np.sin(np.pi*u)*np.cos(np.pi*v); normal=np.c_[-dzdu.ravel(),-dzdv.ravel(),np.ones(q*q)]; normal/=np.linalg.norm(normal,axis=1,keepdims=True)
        faces=[]
        for i in range(q-1):
            for j in range(q-1):
                a=i*q+j; faces += [(a,a+1,a+q),(a+1,a+q+1,a+q)]
        und=np.unique(np.sort(np.asarray([(a,b) for f in faces for a,b in ((f[0],f[1]),(f[1],f[2]),(f[2],f[0]))]),axis=1),axis=0); e=np.c_[np.r_[und[:,0],und[:,1]],np.r_[und[:,1],und[:,0]]].T; delta=p[e[1]]-p[e[0]]; length=np.linalg.norm(delta,axis=1); attr=np.zeros((len(e[0]),edge_dim),np.float32)
        attr[:,0]=length/1000.; attr[:,1:min(4,edge_dim)]=delta[:,:max(0,min(3,edge_dim-1))]/np.maximum(length[:,None],1e-8); 
        if edge_dim>4: attr[:,4]=(normal[e[0]]*normal[e[1]]).sum(1)
        x=np.zeros((q*q,node_dim//2),np.float32); m=np.zeros_like(x); x[:,:3]=p; x[:,3:6]=normal; m[:,:6]=1
        return _tensor(np.c_[x,m]),_tensor(e,torch.long),_tensor(attr)
    a,ae,aa=build(.05); b,be,ba=build(.25); model.eval()
    with torch.no_grad(): delta=float((model(a,ae,aa,cond[:1])-model(b,be,ba,cond[:1])).abs().mean())
    return {"changed_geometry_architecture_only_mean_abs_delta":delta,"changed_geometry_note":"procedural meshes have recomputed positions, normals, topology and edge attributes. This is architecture sensitivity, not AC28 CAD-revision transfer validation."}

def _graph_families(synth: list[dict[str,Any]]) -> list[dict[str,Any]]:
    graphs=graph_inputs(); cond,runs,source=_conditions(); c=graphs[0]["x"].shape[1]*2; e=graphs[0]["edge_attr"].shape[1]; m=MeshGraphProcessor(c,e,cond.shape[1]); rows=[]
    for d in graphs:
        x,ei,ea,g,b=_batched_graph(d,cond); rows.append(_smoke("GNN_MeshGraphNet_style",m,lambda m=m,x=x,ei=ei,ea=ea,g=g,b=b:m(x,ei,ea,g,b),d["source"],"pre-process predictive candidate only with CAD/process inputs and validated outcomes",{"actual_nodes":d["nodes"],"process_runs":runs,"process_source":source,"full_prepared_topology":True,"feature_masks_appended":True,"shared_model_instance_across_budgets":True}))
    # A node-wise three-head candidate for the intended production task. PLY is
    # the supervision source only; actual prediction consumes CAD + process.
    d=max(graphs,key=lambda item:item["nodes"]); mx,mei,mea,mg,mb=_batched_graph(d,cond)
    multi=PhysicsConstrainedRiskGNN(c,e,cond.shape[1],defect_types=3)
    rows.append(_smoke(
        "GNN_CAD_risk_deformation_multitask", multi,
        lambda:multi(mx,mei,mea,mg,mb), d["source"],
        "node-wise pre-process candidate: STEP graph + XLS process -> defect segmentation and deformation fields",
        {"actual_nodes_per_run":d["nodes"], "process_runs":runs, "process_source":source,
         "output_channels":["defect_risk_logit","signed_normal_deformation","deformation_magnitude","graph_defect_type_logits"],
         "defect_type_classes":"three placeholder output slots; class vocabulary requires validated data",
         "hard_constraint":"deformation_magnitude >= abs(signed_normal_deformation)",
         "target_scaffold":"data_processed/paired; currently weak/diagnostic with training_valid=false",
         "ply_role":"training/validation target construction only; not an inference input"}))
    synth.append(_synthetic_checkpoint("cad_risk_deformation_multitask",multi,
                                       lambda:multi(mx,mei,mea,mg,mb)))
    # Exact relabeling test on an actual prepared budget, with every original edge remapped.
    d=graphs[-1]; x,ei,ea,g,b=_batched_graph(d,cond[:1]); n=len(d["x"]); perm=torch.randperm(n); inv=torch.argsort(perm); ep=inv[ei]
    m.eval();
    with torch.no_grad(): equiv=float((m(x,ei,ea,g,b)-m(x[perm],ep,ea,g,b[perm])[inv]).abs().max())
    px,pei,pea,pg,pb=_batched_graph(d,cond)
    with torch.no_grad():
        two=m(px,pei,pea,pg,pb)
        process_delta=float((two[:n]-two[n:]).abs().mean()) if len(cond)>1 else None
    rows[0]["permutation_equivariance_max_abs"] = equiv; rows[0]["actual_process_condition_architecture_only_mean_abs_delta"] = process_delta; rows[0]["process_sensitivity_note"]="Different prepared process values/masks were supplied; untrained response is architecture sensitivity, not process-effect validation."; rows[0].update(_procedural_geometry_probe(m,cond,c,e)); synth.append(_synthetic_checkpoint("meshgraphnet",m,lambda m=m,x=x,ei=ei,ea=ea,g=g,b=b:m(x,ei,ea,g,b)))
    return rows

def _uv_tabular_families(synth: list[dict[str,Any]]) -> list[dict[str,Any]]:
    rows=[]; d=brep_input(); z,ad=_tensor(d["grids"]),_tensor(d["adjacency"],torch.long); m=UVNetFaceModel(z.shape[1]); rows.append(_smoke("UVNet_face_grid_adjacency",m,lambda:m(z,ad),d["source"],"CAD B-Rep pre-process candidate only with validated outcome labels",{"faces":d["faces"],"face_adjacency_directed_edges":d["adjacency_edges"],"trim_and_feature_masks_encoded":True,"adjacency_note":"Used exactly as prepared from shared B-Rep topological edges; no proximity edges were invented."})); synth.append(_synthetic_checkpoint("uvnet",m,lambda:m(z,ad)))
    t=tabular_input(); x=_tensor(np.concatenate([t["x"],t["mask"].astype(np.float32)],1)); m=TabularSurrogate(x.shape[1]); rows.append(_smoke("tabular_surrogate",m,lambda:m(x),t["source"],"pre-process candidate only after independent validated outcomes",{"actual_runs":t["runs"],"feature_masks_appended":True})); synth.append(_synthetic_checkpoint("tabular_surrogate",m,lambda:m(x)))
    return rows

def synthetic_bo() -> dict[str,Any]:
    from sklearn.gaussian_process import GaussianProcessRegressor
    from sklearn.gaussian_process.kernels import Matern,WhiteKernel
    from scipy.stats import norm
    rng=np.random.default_rng(17); x=rng.uniform([0,.2],[1,.8],(28,2)); y=(x[:,0]-.62)**2+1.4*(x[:,1]-.42)**2; gp=GaussianProcessRegressor(Matern(nu=2.5)+WhiteKernel(.001),normalize_y=True,random_state=17).fit(x,y)
    cand=rng.uniform([0,.2],[1,.8],(2000,2)); mean,sd=gp.predict(cand,return_std=True); ei=(y.min()-mean)*norm.cdf((y.min()-mean)/(sd+1e-9))+sd*norm.pdf((y.min()-mean)/(sd+1e-9)); feasible=(cand[:,0]>=.15)&(cand[:,0]<=.85)&(cand[:,1]>=.25)&(cand[:,1]<=.75); pick=int(np.argmax(np.where(feasible,ei,-np.inf)))
    return {"family":"sklearn_GP_constrained_expected_improvement","status":"synthetic_demo","synthetic_only":True,"hard_feasibility_bounds":{"p0":[.15,.85],"p1":[.25,.75]},"candidate":cand[pick].round(5).tolist(),"expected_improvement":float(ei[pick]),"warning":"No AC28 condition recommendation, accuracy, or feasibility inference."}

def main() -> None:
    p=argparse.ArgumentParser(description=__doc__); p.add_argument("--smoke",action="store_true"); p.add_argument("--all",action="store_true"); p.add_argument("--train",type=Path); p.add_argument("--checkpoint",type=Path,default=OUT/"checkpoints/validated_tabular.pt"); args=p.parse_args()
    if not(args.smoke or args.all or args.train): p.error("choose --smoke, --all, or --train")
    OUT.mkdir(parents=True,exist_ok=True); seed_everything(); report={"label_policy":"Current prepared training_valid masks are false. Real supervised training is refused; forward/backward smoke has no accuracy.","models":[],"unavailable":[],"synthetic":[]}
    if args.train:
        # Validation happens before the dimensions are used, giving an actionable refusal.
        cin,cout=validated_contract_dimensions(args.train)
        report["training"]=fit_validated_npz(TabularSurrogate(cin*2,cout),args.train,args.checkpoint)
    if args.smoke or args.all:
        for fun in (_image_families,_voxel_families,_point_families,_graph_families,_uv_tabular_families):
            try: report["models"].extend(fun(report["synthetic"]))
            except (InputContractError,RuntimeError,ValueError) as exc: report["unavailable"].append({"family":fun.__name__,"status":"unavailable_input_or_execution_error","reason":str(exc),"action":"inspect the named prepared schema; no fallback/random input was used"})
    if args.all: report["synthetic"].append(synthetic_bo())
    (OUT/"report.json").write_text(json.dumps(report,indent=2),encoding="utf8"); print(OUT/"report.json")
if __name__=="__main__": main()