"""Focused executable checks for safety-critical baseline behavior."""
import numpy as np
import torch
from pathlib import Path
from tempfile import TemporaryDirectory
from models.architectures import DGCNNSeg, FNO2d, PhysicsConstrainedRiskGNN
from models.architectures import TabularSurrogate
from models.training import DatasetContractError, TrainOnlyScaler, masked_mse, fit_validated_npz, validated_contract_dimensions
from models.data import fixed_physical_normalize, sparse_voxel_input, sparse_batch_collate, split_embedded_units

def test_masked_nan_isolated_from_gradient():
    p=torch.tensor([1.,2.],requires_grad=True); y=torch.tensor([0.,float("nan")]); loss=masked_mse(p,y,torch.tensor([True,False])); loss.backward()
    assert torch.isfinite(p.grad).all() and p.grad[1] == 0

def test_empty_mask_refused():
    try: masked_mse(torch.ones(2),torch.ones(2),torch.zeros(2,dtype=torch.bool))
    except DatasetContractError: return
    raise AssertionError("empty target mask was accepted")

def test_train_scaler_uses_mask_not_zero_imputation():
    s=TrainOnlyScaler().fit(np.array([[10.,999.],[20.,999.]]),np.array([[True,False],[True,False]]))
    assert np.allclose(s.mean,[15.,0.]) and np.allclose(s.transform(np.array([[15.,12.]]),np.array([[True,False]])),[[0.,0.]])

def test_dgcnn_one_point_and_fno_both_sides_backward():
    y=DGCNNSeg(3)(torch.randn(1,1,3)); y.square().mean().backward()
    f=FNO2d(2,modes=8); out=f(torch.randn(1,2,20,12)); out.square().mean().backward()
    assert f.spec[0].weight_pos.grad is not None and f.spec[0].weight_neg.grad is not None

def test_unit_aware_normalization_preserves_common_native_axes_and_curvature():
    x=np.array([[1000.,2000.,-3000.,2.,4.]],np.float32)
    y=fixed_physical_normalize(x,np.array(["x_native","y_native","z_native","curvature_magnitude_per_mm","vertex_area_mm2"]),np.array(["native","native","native","1/mm","mm2"]))
    assert np.allclose(y[0,:3],[1/3,2/3,-1]) and np.isclose(y[0,3],np.log(3)) and np.isclose(y[0,4],4e-6)
    names,units=split_embedded_units(np.array(["상부온도01 [°C]","가열시간 [sec]","위치 [mm]"]))
    scaled=fixed_physical_normalize(np.array([[250.,50.,500.]],np.float32),names,units)
    assert names.tolist()==["상부온도01","가열시간","위치"] and np.allclose(scaled,[[.5,.5,.5]])

def test_training_contract_refuses_current_unlabeled_data_and_accepts_grouped_valid_demo():
    try: validated_contract_dimensions("data_processed/tabular/process_and_geometry.npz")
    except DatasetContractError: pass
    else: raise AssertionError("current unlabeled tabular data was accepted")
    with TemporaryDirectory() as td:
        p=Path(td)/"valid.npz"; x=np.arange(24,dtype=np.float32).reshape(6,4); y=x[:,:1]; mask=np.ones_like(x,bool); valid=np.ones_like(y,bool)
        np.savez(p,x=x,feature_valid=mask,y=y,training_valid=valid,group_id=np.array(["a","a","b","b","c","c"]),label_status=np.array(["externally_validated"]*6))
        result=fit_validated_npz(TabularSurrogate(8,1),p,Path(td)/"ckpt.pt",epochs=1)
        assert result["status"]=="trained" and (Path(td)/"ckpt.pt").exists()

def test_sparse_npz_schema_and_minkowski_coordinate_collation():
    d=sparse_voxel_input(); c,f=sparse_batch_collate(d["coords"],d["features"])
    assert c.dtype==np.int32 and c.shape==(len(d["coords"]),4) and np.all(c[:,0]==0)
    assert np.array_equal(c[:,1:],d["coords"]) and f.shape[0]==c.shape[0] and np.isfinite(f).all()
    try: import MinkowskiEngine as ME  # type: ignore
    except ImportError: return  # backend remains optional on this platform
    from models.architectures import SparseUNetOptional
    model=SparseUNetOptional(f.shape[1]); inp=ME.SparseTensor(features=torch.as_tensor(f),coordinates=torch.as_tensor(c))
    out=model(inp)
    assert torch.equal(torch.unique(out.C,dim=0),torch.unique(inp.C,dim=0)) and torch.isfinite(out.F).all()

def test_physics_risk_gnn_shapes_and_magnitude_constraint():
    model=PhysicsConstrainedRiskGNN(node_dim=4,edge_dim=2,global_dim=3,defect_types=2,width=8,steps=1)
    node=torch.randn(5,4); edge=torch.tensor([[0,1,2,3],[1,2,3,4]]); attr=torch.randn(4,2)
    batch=torch.tensor([0,0,0,1,1]); process=torch.randn(2,3)
    out=model(node,edge,attr,process,batch)
    assert out["risk_logit"].shape==(5,) and out["defect_type_logits"].shape==(2,2)
    assert torch.all(out["deformation_magnitude"] >= out["signed_normal_deformation"].abs())
    sum(value.square().mean() for value in out.values()).backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())