import sys,tempfile,unittest
from pathlib import Path
import numpy as np
from PIL import Image
ROOT=Path(__file__).resolve().parents[1]
sys.path[:0]=[str(ROOT/'code'),str(ROOT/'ui')]
from train_recon import load_dat,Net
from infer_brainmu_lora import metrics,LoRALinear
from multi_gpu import shard_indices
import torch
class CoreTests(unittest.TestCase):
 def test_dat_orientation(self):
  x=np.zeros((41,250,400),np.uint8);x[0,0,0]=1;x[-1,-1,-1]=1
  with tempfile.TemporaryDirectory() as d:
   p=Path(d)/'x.dat';np.packbits(x.reshape(41,-1),axis=1,bitorder='little').tofile(p)
   np.testing.assert_array_equal(load_dat(p),x[:,::-1,:])
 def test_metrics_identity(self):
  m=metrics(Image.new('RGB',(32,32),'gray'),Image.new('RGB',(32,32),'gray'))
  self.assertAlmostEqual(m['ssim'],1);self.assertAlmostEqual(m['psnr_db'],120)
 def test_shards(self):
  s=shard_indices(1000,8);self.assertEqual([len(x) for x in s],[125]*8)
  self.assertEqual(sorted(i for x in s for i in x),list(range(1000)))
 def test_lora_zero(self):
  base=torch.nn.Linear(4,3);m=LoRALinear(base,2,4,0);m.lora_A.data.zero_();m.lora_B.data.zero_();x=torch.ones(1,4)
  torch.testing.assert_close(m(x),base(x))
 def test_frontend(self):
  self.assertEqual(tuple(Net()(torch.zeros(1,41,8,8)).shape),(1,1,8,8))
if __name__=='__main__':unittest.main()