# Copyright (c) OpenMMLab. All rights reserved.
import torch

from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder
from mmedit.models.common import SimpleGatedConvModule


def test_deepfill_enc():
    encoder = DeepFillEncoder()
    x = torch.randn((2, 5, 256, 256))
    outputs = encoder(x)
    assert isinstance(outputs, dict)
    assert 'out' in outputs
    res = outputs['out']
    assert res.shape == (2, 128, 64, 64)
    assert encoder.enc2.stride == (2, 2)
    assert encoder.enc2.out_channels == 64

    encoder = DeepFillEncoder(encoder_type='stage2_conv')
    x = torch.randn((2, 5, 256, 256))
    outputs = encoder(x)
    assert isinstance(outputs, dict)
    assert 'out' in outputs
    res = outputs['out']
    assert res.shape == (2, 128, 64, 64)
    assert encoder.enc2.out_channels == 32
    assert encoder.enc3.out_channels == 64
    assert encoder.enc4.out_channels == 64

    encoder = DeepFillEncoder(encoder_type='stage2_attention')
    x = torch.randn((2, 5, 256, 256))
    outputs = encoder(x)
    assert isinstance(outputs, dict)
    assert 'out' in outputs
    res = outputs['out']
    assert res.shape == (2, 128, 64, 64)
    assert encoder.enc2.out_channels == 32
    assert encoder.enc3.out_channels == 64
    assert encoder.enc4.out_channels == 128
    if torch.cuda.is_available():
        encoder = DeepFillEncoder().cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 128, 64, 64)
        assert encoder.enc2.stride == (2, 2)
        assert encoder.enc2.out_channels == 64

        encoder = DeepFillEncoder(encoder_type='stage2_conv').cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 128, 64, 64)
        assert encoder.enc2.out_channels == 32
        assert encoder.enc3.out_channels == 64
        assert encoder.enc4.out_channels == 64

        encoder = DeepFillEncoder(encoder_type='stage2_attention').cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 128, 64, 64)
        assert encoder.enc2.out_channels == 32
        assert encoder.enc3.out_channels == 64
        assert encoder.enc4.out_channels == 128

        encoder = DeepFillEncoder(
            conv_type='gated_conv', channel_factor=0.75).cuda()
        x = torch.randn((2, 5, 256, 256)).cuda()
        outputs = encoder(x)
        assert isinstance(outputs, dict)
        assert 'out' in outputs
        res = outputs['out']
        assert res.shape == (2, 96, 64, 64)
        assert isinstance(encoder.enc2, SimpleGatedConvModule)
        assert encoder.enc2.conv.stride == (2, 2)
        assert encoder.enc2.conv.out_channels == 48 * 2


def test_deepfill_contextual_attention_neck():
    # TODO: add unittest for contextual attention module
    neck = ContextualAttentionNeck(in_channels=128)
    x = torch.rand((2, 128, 64, 64))
    mask = torch.zeros((2, 1, 64, 64))
    mask[..., 20:100, 23:90] = 1.

    res, offset = neck(x, mask)

    assert res.shape == (2, 128, 64, 64)
    assert offset.shape == (2, 32, 32, 32, 32)

    if torch.cuda.is_available():
        neck.cuda()
        res, offset = neck(x.cuda(), mask.cuda())

        assert res.shape == (2, 128, 64, 64)
        assert offset.shape == (2, 32, 32, 32, 32)

        neck = ContextualAttentionNeck(
            in_channels=128, conv_type='gated_conv').cuda()
        res, offset = neck(x.cuda(), mask.cuda())
        assert res.shape == (2, 128, 64, 64)
        assert offset.shape == (2, 32, 32, 32, 32)
        assert isinstance(neck.conv1, SimpleGatedConvModule)

[(9, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ([], {}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (10, 'torch.randn', 'torch.randn', (['(2, 5, 256, 256)'], {}), False, 'import torch\n'), (19, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ([], {'encoder_type': '"""stage2_conv"""'}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (20, 'torch.randn', 'torch.randn', (['(2, 5, 256, 256)'], {}), False, 'import torch\n'), (30, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ([], {'encoder_type': '"""stage2_attention"""'}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (31, 'torch.randn', 'torch.randn', (['(2, 5, 256, 256)'], {}), False, 'import torch\n'), (40, 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), False, 'import torch\n'), (88, 'mmedit.models.backbones.ContextualAttentionNeck', 'ContextualAttentionNeck', ([], {'in_channels': '(128)'}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (89, 'torch.rand', 'torch.rand', (['(2, 128, 64, 64)'], {}), False, 'import torch\n'), (90, 'torch.zeros', 'torch.zeros', (['(2, 1, 64, 64)'], {}), False, 'import torch\n'), (98, 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), False, 'import torch\n'), (41, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ([], {}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (42, 'torch.randn', 'torch.randn', (['(2, 5, 256, 256)'], {}), False, 'import torch\n'), (51, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ([], {'encoder_type': '"""stage2_conv"""'}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (52, 'torch.randn', 'torch.randn', (['(2, 5, 256, 256)'], {}), False, 'import torch\n'), (62, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ([], {'encoder_type': '"""stage2_attention"""'}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (63, 'torch.randn', 'torch.randn', (['(2, 5, 256, 256)'], {}), False, 'import torch\n'), (73, 'mmedit.models.backbones.DeepFillEncoder', 'DeepFillEncoder', ([], {'conv_type': '"""gated_conv"""', 'channel_factor': '(0.75)'}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n'), (75, 'torch.randn', 'torch.randn', (['(2, 5, 256, 256)'], {}), False, 'import torch\n'), (105, 'mmedit.models.backbones.ContextualAttentionNeck', 'ContextualAttentionNeck', ([], {'in_channels': '(128)', 'conv_type': '"""gated_conv"""'}), False, 'from mmedit.models.backbones import ContextualAttentionNeck, DeepFillEncoder\n')]