import pytest
import torch

from mmedit.models.builder import build_component
from mmedit.models.components.discriminators.light_cnn import MaxFeature


def test_max_feature():
    # cpu
    conv2d = MaxFeature(16, 16, filter_type='conv2d')
    x1 = torch.rand(3, 16, 16, 16)
    y1 = conv2d(x1)
    assert y1.shape == (3, 16, 16, 16)
    linear = MaxFeature(16, 16, filter_type='linear')
    x2 = torch.rand(3, 16)
    y2 = linear(x2)
    assert y2.shape == (3, 16)
    # gpu
    if torch.cuda.is_available():
        x1 = x1.cuda()
        x2 = x2.cuda()
        conv2d = conv2d.cuda()
        linear = linear.cuda()
        y1 = conv2d(x1)
        assert y1.shape == (3, 16, 16, 16)
        y2 = linear(x2)
        assert y2.shape == (3, 16)
    # filter_type should be conv2d or linear
    with pytest.raises(ValueError):
        MaxFeature(12, 12, filter_type='conv1d')


def test_light_cnn():
    cfg = dict(type='LightCNN', in_channels=3)
    net = build_component(cfg)
    net.init_weights(pretrained=None)
    # cpu
    inputs = torch.rand((2, 3, 128, 128))
    output = net(inputs)
    assert output.shape == (2, 1)
    # gpu
    if torch.cuda.is_available():
        net.init_weights(pretrained=None)
        net = net.cuda()
        output = net(inputs.cuda())
        assert output.shape == (2, 1)

    # pretrained should be str or None
    with pytest.raises(TypeError):
        net.init_weights(pretrained=[1])

[((200, 240), 'mmedit.models.components.discriminators.light_cnn.MaxFeature', 'MaxFeature', (['(16)', '(16)'], {'filter_type': '"""conv2d"""'}), "(16, 16, filter_type='conv2d')\n", (210, 240), False, 'from mmedit.models.components.discriminators.light_cnn import MaxFeature\n'), ((250, 275), 'torch.rand', 'torch.rand', (['(3)', '(16)', '(16)', '(16)'], {}), '(3, 16, 16, 16)\n', (260, 275), False, 'import torch\n'), ((348, 388), 'mmedit.models.components.discriminators.light_cnn.MaxFeature', 'MaxFeature', (['(16)', '(16)'], {'filter_type': '"""linear"""'}), "(16, 16, filter_type='linear')\n", (358, 388), False, 'from mmedit.models.components.discriminators.light_cnn import MaxFeature\n'), ((398, 415), 'torch.rand', 'torch.rand', (['(3)', '(16)'], {}), '(3, 16)\n', (408, 415), False, 'import torch\n'), ((484, 509), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (507, 509), False, 'import torch\n'), ((956, 976), 'mmedit.models.builder.build_component', 'build_component', (['cfg'], {}), '(cfg)\n', (971, 976), False, 'from mmedit.models.builder import build_component\n'), ((1038, 1066), 'torch.rand', 'torch.rand', (['(2, 3, 128, 128)'], {}), '((2, 3, 128, 128))\n', (1048, 1066), False, 'import torch\n'), ((1143, 1168), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1166, 1168), False, 'import torch\n'), ((799, 824), 'pytest.raises', 'pytest.raises', (['ValueError'], {}), '(ValueError)\n', (812, 824), False, 'import pytest\n'), ((834, 874), 'mmedit.models.components.discriminators.light_cnn.MaxFeature', 'MaxFeature', (['(12)', '(12)'], {'filter_type': '"""conv1d"""'}), "(12, 12, filter_type='conv1d')\n", (844, 874), False, 'from mmedit.models.components.discriminators.light_cnn import MaxFeature\n'), ((1360, 1384), 'pytest.raises', 'pytest.raises', (['TypeError'], {}), '(TypeError)\n', (1373, 1384), False, 'import pytest\n')]