import unittest,importlib.util
AVAILABLE=importlib.util.find_spec('torch') is not None
@unittest.skipUnless(AVAILABLE,'PyTorch is unavailable; gradient/training tests NOT executed')
class TorchTests(unittest.TestCase):
 def test_shapes_and_gradient(self):
  import torch
  from irstd.models import TinyUNet,segmentation_loss
  model=TinyUNet();x=torch.randn(1,1,17,19);y=torch.zeros_like(x);y[:,:,8,9]=1
  for name in ['bce','dice','bce_dice','focal','iou']:
   model.zero_grad();z=model(x);self.assertEqual(z.shape,x.shape)
   loss=segmentation_loss(z,y,name);loss.backward();self.assertTrue(torch.isfinite(loss))
   self.assertTrue(all(torch.isfinite(p.grad).all() for p in model.parameters()))
 def test_gate_gradients(self):
  import torch
  from irstd.models import gate_loss
  s=torch.tensor([0.,0.],requires_grad=True)
  loss,upper=gate_loss(s,torch.tensor([1.,0.]),torch.tensor([4.,1.]),1.,.1)
  loss.backward();self.assertLess(s.grad[0],0);self.assertGreater(s.grad[1],0);self.assertAlmostEqual(upper.item(),1.)