import unittest

import torch
from parameterized import parameterized

from src.constructor import create_backbone
from src.models.backbones.utils import list_models
from .test_segmentation import example_backbones


def inp(bsize, in_ch, w, h):
    return torch.ones(bsize, in_ch, w, h)


class TestBackboneCorrectness(unittest.TestCase):
    def setUp(self) -> None:
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    @parameterized.expand(list_models(module='vision_transformer', exclude_filters=''))
    def test_vit_torchscript_conversion(self, backbone_name):
        model = create_backbone(backbone_name, img_size=self.input.shape[2]).to(self.device).eval()
        with torch.no_grad():
            torch.jit.trace(model, self.input)
        torch.cuda.empty_cache()

    @parameterized.expand(list_models(module='coat', exclude_filters=''))
    def test_coat_torchscript_conversion(self, backbone_name):
        model = create_backbone(backbone_name, img_size=self.input.shape[2]).to(self.device).eval()
        with torch.no_grad():
            torch.jit.trace(model, self.input)
        torch.cuda.empty_cache()

    @parameterized.expand(list_models(module='swin_transformer', exclude_filters=''))
    def test_swin_torchscript_conversion(self, backbone_name):
        model = create_backbone(backbone_name).to(self.device).eval()
        input = torch.rand(2, 3, *model.img_size, device=self.device)
        with torch.no_grad():
            torch.jit.trace(model, input)
        torch.cuda.empty_cache()

[((257, 287), 'torch.ones', 'torch.ones', (['bsize', 'in_ch', 'w', 'h'], {}), '(bsize, in_ch, w, h)\n', (267, 287), False, 'import torch\n'), ((790, 814), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ([], {}), '()\n', (812, 814), False, 'import torch\n'), ((481, 541), 'src.models.backbones.utils.list_models', 'list_models', ([], {'module': '"""vision_transformer"""', 'exclude_filters': '""""""'}), "(module='vision_transformer', exclude_filters='')\n", (492, 541), False, 'from src.models.backbones.utils import list_models\n'), ((1138, 1162), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ([], {}), '()\n', (1160, 1162), False, 'import torch\n'), ((842, 888), 'src.models.backbones.utils.list_models', 'list_models', ([], {'module': '"""coat"""', 'exclude_filters': '""""""'}), "(module='coat', exclude_filters='')\n", (853, 888), False, 'from src.models.backbones.utils import list_models\n'), ((1399, 1452), 'torch.rand', 'torch.rand', (['(2)', '(3)', '*model.img_size'], {'device': 'self.device'}), '(2, 3, *model.img_size, device=self.device)\n', (1409, 1452), False, 'import torch\n'), ((1533, 1557), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ([], {}), '()\n', (1555, 1557), False, 'import torch\n'), ((1190, 1248), 'src.models.backbones.utils.list_models', 'list_models', ([], {'module': '"""swin_transformer"""', 'exclude_filters': '""""""'}), "(module='swin_transformer', exclude_filters='')\n", (1201, 1248), False, 'from src.models.backbones.utils import list_models\n'), ((718, 733), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (731, 733), False, 'import torch\n'), ((747, 781), 'torch.jit.trace', 'torch.jit.trace', (['model', 'self.input'], {}), '(model, self.input)\n', (762, 781), False, 'import torch\n'), ((1066, 1081), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (1079, 1081), False, 'import torch\n'), ((1095, 1129), 'torch.jit.trace', 'torch.jit.trace', (['model', 'self.input'], {}), '(model, self.input)\n', (1110, 1129), False, 'import torch\n'), ((1466, 1481), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (1479, 1481), False, 'import torch\n'), ((1495, 1524), 'torch.jit.trace', 'torch.jit.trace', (['model', 'input'], {}), '(model, input)\n', (1510, 1524), False, 'import torch\n'), ((416, 441), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (439, 441), False, 'import torch\n'), ((621, 681), 'src.constructor.create_backbone', 'create_backbone', (['backbone_name'], {'img_size': 'self.input.shape[2]'}), '(backbone_name, img_size=self.input.shape[2])\n', (636, 681), False, 'from src.constructor import create_backbone\n'), ((969, 1029), 'src.constructor.create_backbone', 'create_backbone', (['backbone_name'], {'img_size': 'self.input.shape[2]'}), '(backbone_name, img_size=self.input.shape[2])\n', (984, 1029), False, 'from src.constructor import create_backbone\n'), ((1329, 1359), 'src.constructor.create_backbone', 'create_backbone', (['backbone_name'], {}), '(backbone_name)\n', (1344, 1359), False, 'from src.constructor import create_backbone\n')]