# Copyright (c) OpenMMLab. All rights reserved.
from copy import deepcopy

import numpy as np
import pytest
import torch
import torch.nn as nn

from mmgen.models.gans import ProgressiveGrowingGAN


class TestPGGAN:

    @classmethod
    def setup_class(cls):
        cls.generator_cfg = dict(
            type='PGGANGenerator',
            noise_size=8,
            out_scale=16,
            base_channels=32,
            max_channels=32)

        cls.discriminator_cfg = dict(
            type='PGGANDiscriminator', in_scale=16, label_size=0)

        cls.gan_loss = dict(type='GANLoss', gan_type='vanilla')
        cls.disc_auxiliary_loss = [
            dict(
                type='DiscShiftLoss',
                loss_weight=0.5,
                data_info=dict(pred='disc_pred_fake')),
            dict(
                type='DiscShiftLoss',
                loss_weight=0.5,
                data_info=dict(pred='disc_pred_real'))
        ]
        cls.train_cfg = dict(
            use_ema=True,
            nkimgs_per_scale={
                '4': 0.004,
                '8': 0.008,
                '16': 0.016
            },
            optimizer_cfg=dict(
                generator=dict(type='Adam', lr=0.0002, betas=(0.5, 0.999)),
                discriminator=dict(type='Adam', lr=0.0002,
                                   betas=(0.5, 0.999))),
            g_lr_base=0.0001,
            d_lr_base=0.0001,
            g_lr_schedule={'16': 0.00005})

    def test_pggan_cpu(self):
        # test default config
        pggan = ProgressiveGrowingGAN(
            self.generator_cfg,
            self.discriminator_cfg,
            self.gan_loss,
            disc_auxiliary_loss=self.disc_auxiliary_loss,
            train_cfg=self.train_cfg)

        data_batch = dict(real_img=torch.randn(3, 3, 16, 16))

        for iter_num in range(6):
            outputs = pggan.train_step(
                data_batch,
                None,
                running_status=dict(iteration=iter_num, batch_size=3))
            results = outputs['results']
            if iter_num == 1:
                assert results['fake_imgs'].shape == (3, 3, 4, 4)
            elif iter_num == 2:
                assert results['fake_imgs'].shape == (3, 3, 8, 8)
                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
            elif iter_num == 3:
                assert results['fake_imgs'].shape == (3, 3, 8, 8)
                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
                assert np.isclose(
                    pggan.optimizer['generator'].defaults['lr'],
                    0.0001,
                    atol=1e-8)
            elif iter_num == 5:
                assert results['fake_imgs'].shape == (3, 3, 16, 16)
                assert np.isclose(pggan._actual_nkimgs[-1], 0.012, atol=1e-8)
                assert np.isclose(
                    pggan.optimizer['generator'].defaults['lr'],
                    0.00005,
                    atol=1e-8)

        # test sample from noise
        outputs = pggan.sample_from_noise(None, num_batches=2)
        assert outputs.shape == (4, 3, 16, 16)

        outputs = pggan.sample_from_noise(
            None,
            num_batches=2,
            return_noise=True,
            transition_weight=0.2,
            sample_model='ema')
        assert outputs['fake_img'].shape == (2, 3, 16, 16)

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

        with pytest.raises(RuntimeError):
            data_batch = dict(real_img=torch.randn(3, 3, 4, 32))
            _ = pggan.train_step(
                data_batch, None, running_status=dict(iteration=5))

        # test customized config
        train_cfg_ = deepcopy(self.train_cfg)
        train_cfg_['use_ema'] = False
        train_cfg_['interp_real_cfg'] = dict(
            mode='bilinear', align_corners=False)
        train_cfg_['interp_real_cfg'] = dict(
            mode='bilinear', align_corners=False)
        train_cfg_['reset_optim_for_new_scale'] = False
        pggan = ProgressiveGrowingGAN(
            self.generator_cfg,
            self.discriminator_cfg,
            self.gan_loss,
            disc_auxiliary_loss=None,
            train_cfg=train_cfg_)

        data_batch = dict(real_img=torch.randn(3, 3, 16, 16))

        outputs = pggan.train_step(
            data_batch, None, running_status=dict(iteration=0, batch_size=3))
        results = outputs['results']
        assert results['fake_imgs'].shape == (3, 3, 4, 4)
        assert not pggan.with_gen_auxiliary_loss
        assert not pggan.with_disc_auxiliary_loss
        assert not pggan.use_ema

        data_batch = dict(real_img=torch.randn(3, 3, 16, 16))

        for iter_num in range(1, 3):
            outputs = pggan.train_step(
                data_batch,
                None,
                running_status=dict(iteration=iter_num, batch_size=3))
            results = outputs['results']
            if iter_num == 1:
                assert results['fake_imgs'].shape == (3, 3, 4, 4)
            elif iter_num == 2:
                assert results['fake_imgs'].shape == (3, 3, 8, 8)
                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)

        train_cfg_ = deepcopy(self.train_cfg)
        train_cfg_['optimizer_cfg'] = dict(
            generator=dict(type='Adam', lr=0.0002, betas=(0.5, 0.999)))
        pggan = ProgressiveGrowingGAN(
            self.generator_cfg,
            None,
            None,
            disc_auxiliary_loss=dict(
                type='DiscShiftLoss',
                loss_weight=0.5,
                data_info=dict(pred='disc_pred_fake')),
            gen_auxiliary_loss=dict(type='GeneratorPathRegularizer'),
            train_cfg=train_cfg_)
        assert pggan.with_gen_auxiliary_loss
        assert isinstance(pggan.disc_auxiliary_losses, nn.ModuleList)
        assert pggan.gan_loss is None
        assert pggan.discriminator is None

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

        data_batch = dict(real_img=torch.randn(3, 3, 32, 32).cuda())

        for iter_num in range(6):
            outputs = pggan.train_step(
                data_batch,
                None,
                running_status=dict(iteration=iter_num, batch_size=3))
            results = outputs['results']
            if iter_num == 1:
                assert results['fake_imgs'].shape == (3, 3, 4, 4)
            elif iter_num == 2:
                assert results['fake_imgs'].shape == (3, 3, 8, 8)
                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
            elif iter_num == 3:
                assert results['fake_imgs'].shape == (3, 3, 8, 8)
                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
            elif iter_num == 5:
                assert results['fake_imgs'].shape == (3, 3, 16, 16)
                assert np.isclose(pggan._actual_nkimgs[-1], 0.012, atol=1e-8)
