import json,tempfile,hashlib
from pathlib import Path
import torch
from semantic_model import SemanticColorizer,save_semantic,load_semantic
from objectives import color_loss
torch.set_num_threads(2);torch.manual_seed(609)
results=[]
for head in ['dense','palette']:
 m=SemanticColorizer(head=head,pretrained=True).eval()
 count=sum(p.numel() for p in m.parameters());assert count<4_000_000
 for shape in [(2,1,64,80),(1,1,173,241),(1,1,256,256)]:
  with torch.no_grad():y=m(torch.zeros(shape))
  assert y.shape==(shape[0],2,shape[2],shape[3]) and torch.isfinite(y).all()
 x=torch.rand(2,1,64,64)*2-1;t=torch.zeros(2,2,64,64);t[:,0]=12;t[:,1]=18
 opt=torch.optim.AdamW([p for n,p in m.named_parameters() if not n.startswith('encoder.')],lr=.001)
 with torch.no_grad():initial=float(color_loss(m(x),t)[0])
 for _ in range(16):
  opt.zero_grad();loss,_=color_loss(m(x),t);loss.backward()
  assert all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)
  opt.step()
 with torch.no_grad():final=float(color_loss(m(x),t)[0]);ref=m(x)
 assert final<initial,(head,initial,final)
 with tempfile.TemporaryDirectory() as tmp:
  save_semantic(m,tmp);restored=load_semantic(tmp)
  with torch.no_grad():torch.testing.assert_close(restored(x),ref)
 results.append({'head':head,'parameters':count,'initial_loss':initial,'final_loss':final,'roundtrip':'passed','shape_tests':'passed','finite_backward':'passed'})
print('SMOKE_RESULTS',json.dumps(results),flush=True)
Path('smoke_results.json').write_text(json.dumps(results,indent=2))