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()

[((12, 11, 12, 41), 'torch.ones', 'torch.ones', ({(12, 22, 12, 27): 'bsize', (12, 29, 12, 34): 'in_ch', (12, 36, 12, 37): 'w', (12, 39, 12, 40): 'h'}, {}), '(bsize, in_ch, w, h)', False, 'import torch\n'), ((24, 8, 24, 32), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ({}, {}), '()', False, 'import torch\n'), ((19, 26, 19, 86), 'src.models.backbones.utils.list_models', 'list_models', (), '', False, 'from src.models.backbones.utils import list_models\n'), ((31, 8, 31, 32), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ({}, {}), '()', False, 'import torch\n'), ((26, 26, 26, 72), 'src.models.backbones.utils.list_models', 'list_models', (), '', False, 'from src.models.backbones.utils import list_models\n'), ((36, 16, 36, 69), 'torch.rand', 'torch.rand', (), '', False, 'import torch\n'), ((39, 8, 39, 32), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ({}, {}), '()', False, 'import torch\n'), ((33, 26, 33, 84), 'src.models.backbones.utils.list_models', 'list_models', (), '', False, 'from src.models.backbones.utils import list_models\n'), ((22, 13, 22, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((23, 12, 23, 46), 'torch.jit.trace', 'torch.jit.trace', ({(23, 28, 23, 33): 'model', (23, 35, 23, 45): 'self.input'}, {}), '(model, self.input)', False, 'import torch\n'), ((29, 13, 29, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((30, 12, 30, 46), 'torch.jit.trace', 'torch.jit.trace', ({(30, 28, 30, 33): 'model', (30, 35, 30, 45): 'self.input'}, {}), '(model, self.input)', False, 'import torch\n'), ((37, 13, 37, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((38, 12, 38, 41), 'torch.jit.trace', 'torch.jit.trace', ({(38, 28, 38, 33): 'model', (38, 35, 38, 40): 'input'}, {}), '(model, input)', False, 'import torch\n'), ((17, 47, 17, 72), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((21, 16, 21, 76), 'src.constructor.create_backbone', 'create_backbone', (), '', False, 'from src.constructor import create_backbone\n'), ((28, 16, 28, 76), 'src.constructor.create_backbone', 'create_backbone', (), '', False, 'from src.constructor import create_backbone\n'), ((35, 16, 35, 46), 'src.constructor.create_backbone', 'create_backbone', ({(35, 32, 35, 45): 'backbone_name'}, {}), '(backbone_name)', False, 'from src.constructor import create_backbone\n')]