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

[((10, 13, 10, 53), 'mmedit.models.components.discriminators.light_cnn.MaxFeature', 'MaxFeature', (), '', False, 'from mmedit.models.components.discriminators.light_cnn import MaxFeature\n'), ((11, 9, 11, 34), 'torch.rand', 'torch.rand', ({(11, 20, 11, 21): '3', (11, 23, 11, 25): '16', (11, 27, 11, 29): '16', (11, 31, 11, 33): '16'}, {}), '(3, 16, 16, 16)', False, 'import torch\n'), ((14, 13, 14, 53), 'mmedit.models.components.discriminators.light_cnn.MaxFeature', 'MaxFeature', (), '', False, 'from mmedit.models.components.discriminators.light_cnn import MaxFeature\n'), ((15, 9, 15, 26), 'torch.rand', 'torch.rand', ({(15, 20, 15, 21): '3', (15, 23, 15, 25): '16'}, {}), '(3, 16)', False, 'import torch\n'), ((19, 7, 19, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((35, 10, 35, 30), 'mmedit.models.builder.build_component', 'build_component', ({(35, 26, 35, 29): 'cfg'}, {}), '(cfg)', False, 'from mmedit.models.builder import build_component\n'), ((38, 13, 38, 41), 'torch.rand', 'torch.rand', ({(38, 24, 38, 40): '(2, 3, 128, 128)'}, {}), '((2, 3, 128, 128))', False, 'import torch\n'), ((42, 7, 42, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((29, 9, 29, 34), 'pytest.raises', 'pytest.raises', ({(29, 23, 29, 33): 'ValueError'}, {}), '(ValueError)', False, 'import pytest\n'), ((30, 8, 30, 48), 'mmedit.models.components.discriminators.light_cnn.MaxFeature', 'MaxFeature', (), '', False, 'from mmedit.models.components.discriminators.light_cnn import MaxFeature\n'), ((49, 9, 49, 33), 'pytest.raises', 'pytest.raises', ({(49, 23, 49, 32): 'TypeError'}, {}), '(TypeError)', False, 'import pytest\n')]