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

from mmgen.models.gans import StaticUnconditionalGAN


class TestWGANGP:

    @classmethod
    def setup_class(cls):
        cls.generator_cfg = dict(
            type='WGANGPGenerator', noise_size=128, out_scale=128)

        cls.discriminator_cfg = dict(
            type='WGANGPDiscriminator',
            in_channel=3,
            in_scale=128,
            conv_module_cfg=dict(
                conv_cfg=None,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=True,
                act_cfg=dict(type='LeakyReLU', negative_slope=0.2),
                norm_cfg=dict(type='GN'),
                order=('conv', 'norm', 'act')))

        cls.gan_loss = dict(type='GANLoss', gan_type='wgan')
        cls.disc_auxiliary_loss = dict(
            type='GradientPenaltyLoss',
            loss_weight=10,
            norm_mode='pixel',
            data_info=dict(
                discriminator='disc',
                real_data='real_imgs',
                fake_data='fake_imgs'))
        cls.train_cfg = None

    def test_wgangp_cpu(self):
        # test default config
        wgangp = StaticUnconditionalGAN(
            self.generator_cfg,
            self.discriminator_cfg,
            self.gan_loss,
            disc_auxiliary_loss=self.disc_auxiliary_loss,
            train_cfg=self.train_cfg)

        # test sample from noise
        outputs = wgangp.sample_from_noise(None, num_batches=2)
        assert outputs.shape == (2, 3, 128, 128)

        outputs = wgangp.sample_from_noise(
            None, num_batches=2, return_noise=True, sample_model='orig')
        assert outputs['fake_img'].shape == (2, 3, 128, 128)

        # test train step
        data = torch.randn((2, 3, 128, 128))
        data_input = dict(real_img=data)
        optimizer_g = torch.optim.SGD(wgangp.generator.parameters(), lr=0.01)
        optimizer_d = torch.optim.SGD(
            wgangp.discriminator.parameters(), lr=0.01)
        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)

        model_outputs = wgangp.train_step(data_input, optim_dict)
        assert 'results' in model_outputs
        assert 'log_vars' in model_outputs
        assert model_outputs['num_samples'] == 2

    @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
    def test_wgangp_cuda(self):
        # test default config
        wgangp = StaticUnconditionalGAN(
            self.generator_cfg,
            self.discriminator_cfg,
            self.gan_loss,
            disc_auxiliary_loss=self.disc_auxiliary_loss,
            train_cfg=self.train_cfg).cuda()

        # test sample from noise
        outputs = wgangp.sample_from_noise(None, num_batches=2)
        assert outputs.shape == (2, 3, 128, 128)

        outputs = wgangp.sample_from_noise(
            None, num_batches=2, return_noise=True, sample_model='orig')
        assert outputs['fake_img'].shape == (2, 3, 128, 128)

        # test train step
        data = torch.randn((2, 3, 128, 128)).cuda()
        data_input = dict(real_img=data)
        optimizer_g = torch.optim.SGD(wgangp.generator.parameters(), lr=0.01)
        optimizer_d = torch.optim.SGD(
            wgangp.discriminator.parameters(), lr=0.01)
        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)

        model_outputs = wgangp.train_step(data_input, optim_dict)
        assert 'results' in model_outputs
        assert 'log_vars' in model_outputs
        assert model_outputs['num_samples'] == 2