import torch

from gfpgan.archs.stylegan2_clean_arch import StyleGAN2GeneratorClean


def test_stylegan2generatorclean():
    """Test arch: StyleGAN2GeneratorClean."""

    # model init and forward (gpu)
    if torch.cuda.is_available():
        net = StyleGAN2GeneratorClean(
            out_size=32, num_style_feat=512, num_mlp=8, channel_multiplier=1, narrow=0.5).cuda().eval()
        style = torch.rand((1, 512), dtype=torch.float32).cuda()
        output = net([style], input_is_latent=False)
        assert output[0].shape == (1, 3, 32, 32)
        assert output[1] is None

        # -------------------- with return_latents ----------------------- #
        output = net([style], input_is_latent=True, return_latents=True)
        assert output[0].shape == (1, 3, 32, 32)
        assert len(output[1]) == 1
        # check latent
        assert output[1][0].shape == (8, 512)

        # -------------------- with randomize_noise = False ----------------------- #
        output = net([style], randomize_noise=False)
        assert output[0].shape == (1, 3, 32, 32)
        assert output[1] is None

        # -------------------- with truncation = 0.5 and mixing----------------------- #
        output = net([style, style], truncation=0.5, truncation_latent=style)
        assert output[0].shape == (1, 3, 32, 32)
        assert output[1] is None

        # ------------------ test make_noise ----------------------- #
        out = net.make_noise()
        assert len(out) == 7
        assert out[0].shape == (1, 1, 4, 4)
        assert out[1].shape == (1, 1, 8, 8)
        assert out[2].shape == (1, 1, 8, 8)
        assert out[3].shape == (1, 1, 16, 16)
        assert out[4].shape == (1, 1, 16, 16)
        assert out[5].shape == (1, 1, 32, 32)
        assert out[6].shape == (1, 1, 32, 32)

        # ------------------ test get_latent ----------------------- #
        out = net.get_latent(style)
        assert out.shape == (1, 512)

        # ------------------ test mean_latent ----------------------- #
        out = net.mean_latent(2)
        assert out.shape == (1, 512)