import copy,json,sys,tempfile,unittest
from pathlib import Path
import numpy as np
import torch
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'code'))
from project_config import load_config
from train_recon import Net,load_dat
class ConfigTests(unittest.TestCase):
 def test_checkpoint_compatible_structure(self):
  model=Net();ch=[41,64,128,64,16,1];parts=[]
  for i in range(5):
   parts.append(torch.nn.Conv2d(ch[i],ch[i+1],3,1,1))
   if i<4:parts.append(torch.nn.ReLU(True))
  baseline=torch.nn.Sequential(*parts);baseline.load_state_dict(model.seq.state_dict())
  x=torch.randn(1,41,8,8)
  torch.testing.assert_close(model(x),baseline(x),rtol=0,atol=0)
 def test_configuration_changes_model_and_decoder(self):
  c=load_config();c['input'].update(frames=3,height=4,width=8,packed_bytes=12,flip_height=False);c['frontend']['channels']=[3,5,1]
  with tempfile.TemporaryDirectory() as d:
   p=Path(d)/'config.json';p.write_text(json.dumps(c));raw=np.zeros((3,4,8),np.uint8);raw[0,0,0]=1;f=Path(d)/'x.dat';np.packbits(raw.reshape(3,-1),axis=1,bitorder='little').tofile(f)
   np.testing.assert_array_equal(load_dat(f,p),raw)
   self.assertEqual(Net(p)(torch.zeros(1,3,4,8)).shape,(1,1,4,8))
 def test_reject_invalid_config(self):
  c=load_config();c['input']['packed_bytes']=1
  with tempfile.TemporaryDirectory() as d:
   p=Path(d)/'config.json';p.write_text(json.dumps(c))
   with self.assertRaises(ValueError):load_config(p)