# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is licensed under a Creative Commons
# Attribution-NonCommercial-ShareAlike 4.0 International License.
# You should have received a copy of the license along with this
# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/

"""Generate random images using the techniques described in the paper
"Elucidating the Design Space of Diffusion-Based Generative Models"."""

import os
import re
import click
import tqdm
import pickle
import numpy as np
import torch
import PIL.Image
import dnnlib
import discriminator_lib
from torch_utils import distributed as dist
from torchvision.utils import save_image, make_grid

#----------------------------------------------------------------------------
# Proposed EDM sampler (Algorithm 2).
def edm_sampler(
        net, latents, class_labels=None, randn_like=torch.randn_like,
        num_steps=18, sigma_min=0.002, sigma_max=80, rho=7,
        S_churn=0, S_min=0, S_max=float('inf'), S_noise=1,
        discriminator=None, dg_weight_1st_order=2.0, dg_weight_2nd_order=0.0, S_clip_min=0, S_clip_max=float('inf'), dg_logit_max=1.0,
):
    # Adjust noise levels based on what's supported by the network.
    sigma_min = max(sigma_min, net.sigma_min)
    sigma_max = min(sigma_max, net.sigma_max)

    # Time step discretization.
    step_indices = torch.arange(num_steps, dtype=torch.float64, device=latents.device)
    t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (
                sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
    t_steps = torch.cat([net.round_sigma(t_steps), torch.zeros_like(t_steps[:1])])  # t_N = 0

    # Main sampling loop.
    x_next = latents.to(torch.float64) * t_steps[0]
    for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): # 0, ..., N-1
        x_cur = x_next

        # Increase noise temporarily.
        gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0
        t_hat = net.round_sigma(t_cur + gamma * t_cur)
        x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * randn_like(x_cur)
        # Euler step.
        denoised = net(x_hat, t_hat, class_labels).to(torch.float64)
        d_cur = (x_hat - denoised) / t_hat
        ## DG correction
        if discriminator is not None and dg_weight_1st_order != 0.:
            discriminator_guidance, pred = discriminator_lib.get_grad_log_ratio(discriminator, x_hat, t_hat,
                                                                                     net.img_resolution, S_clip_min,
                                                                                     S_clip_max, class_labels, logit=True)
            # if boosting:
            #     if i % period_weight == 0:
            #         discriminator_guidance[log_ratio < 0.] *= 2.
            candidate = (pred < dg_logit_max).squeeze().bool()
            d_cur[candidate] += dg_weight_1st_order * (discriminator_guidance[candidate] / t_hat) 
            denoised[candidate] -= dg_weight_1st_order * discriminator_guidance[candidate]
        x_next = x_hat + (t_next - t_hat) * d_cur

        # Apply 2nd order correction.
        if i < num_steps - 1:
            denoised = net(x_next, t_next, class_labels).to(torch.float64)
            d_prime = (x_next - denoised) / t_next
            ## DG correction
            if discriminator is not None and dg_weight_2nd_order != 0.:
                discriminator_guidance, pred = discriminator_lib.get_grad_log_ratio(discriminator, x_next, t_next,
                                                                              net.img_resolution, S_clip_min,
                                                                              S_clip_max, class_labels, logit=True)
                candidate = (pred < dg_logit_max).squeeze().bool()
                d_prime[candidate] += dg_weight_2nd_order * (discriminator_guidance[candidate] / t_next)
                denoised[candidate] -= dg_weight_2nd_order * discriminator_guidance[candidate]
            x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime)

    return x_next

# ----------------------------------------------------------------------------
# Wrapper for torch.Generator that allows specifying a different random seed
# for each sample in a minibatch.

class StackedRandomGenerator:
    def __init__(self, device, seeds):
        super().__init__()
        self.generators = [torch.Generator(device).manual_seed(int(seed) % (1 << 32)) for seed in seeds]

    def randn(self, size, **kwargs):
        assert size[0] == len(self.generators)
        return torch.stack([torch.randn(size[1:], generator=gen, **kwargs) for gen in self.generators])

    def randn_like(self, input):
        return self.randn(input.shape, dtype=input.dtype, layout=input.layout, device=input.device)

    def randint(self, *args, size, **kwargs):
        assert size[0] == len(self.generators)
        return torch.stack([torch.randint(*args, size=size[1:], generator=gen, **kwargs) for gen in self.generators])


# ----------------------------------------------------------------------------
# Parse a comma separated list of numbers or ranges and return a list of ints.
# Example: '1,2,5-10' returns [1, 2, 5, 6, 7, 8, 9, 10]

def parse_int_list(s):
    if isinstance(s, list): return s
    ranges = []
    range_re = re.compile(r'^(\d+)-(\d+)$')
    for p in s.split(','):
        m = range_re.match(p)
        if m:
            ranges.extend(range(int(m.group(1)), int(m.group(2)) + 1))
        else:
            ranges.append(int(p))
    return ranges


# ----------------------------------------------------------------------------

@click.command()
@click.option('--network', 'network_pkl',  help='Network pickle filename', metavar='PATH|URL',                      type=str, required=True)
@click.option('--outdir',                  help='Where to save the output images', metavar='DIR',                   type=str, required=True)
@click.option('--seeds',                   help='Random seeds (e.g. 1,2,5-10)', metavar='LIST',                     type=parse_int_list, default='0-63', show_default=True)
@click.option('--subdirs',                 help='Create subdirectory for every 1000 seeds',                         is_flag=True)
@click.option('--class', 'class_idx',      help='Class label  [default: random]', metavar='INT',                    type=click.IntRange(min=0), default=None)
@click.option('--batch', 'max_batch_size', help='Maximum batch size', metavar='INT',                                type=click.IntRange(min=1), default=64, show_default=True)

@click.option('--steps', 'num_steps',      help='Number of sampling steps', metavar='INT',                          type=click.IntRange(min=1), default=18, show_default=True)
@click.option('--sigma_min',               help='Lowest noise level  [default: varies]', metavar='FLOAT',           type=click.FloatRange(min=0, min_open=True))
@click.option('--sigma_max',               help='Highest noise level  [default: varies]', metavar='FLOAT',          type=click.FloatRange(min=0, min_open=True))
@click.option('--rho',                     help='Time step exponent', metavar='FLOAT',                              type=click.FloatRange(min=0, min_open=True), default=7, show_default=True)
@click.option('--S_churn', 'S_churn',      help='Stochasticity strength', metavar='FLOAT',                          type=click.FloatRange(min=0), default=0, show_default=True)
@click.option('--S_min', 'S_min',          help='Stoch. min noise level', metavar='FLOAT',                          type=click.FloatRange(min=0), default=0, show_default=True)
@click.option('--S_max', 'S_max',          help='Stoch. max noise level', metavar='FLOAT',                          type=click.FloatRange(min=0), default='inf', show_default=True)
@click.option('--S_noise', 'S_noise',      help='Stoch. noise inflation', metavar='FLOAT',                          type=float, default=1, show_default=True)

@click.option('--solver',                  help='Ablate ODE solver', metavar='euler|heun',                          type=click.Choice(['euler', 'heun']))
@click.option('--disc', 'discretization',  help='Ablate time step discretization {t_i}', metavar='vp|ve|iddpm|edm', type=click.Choice(['vp', 've', 'iddpm', 'edm']))
@click.option('--schedule',                help='Ablate noise schedule sigma(t)', metavar='vp|ve|linear',           type=click.Choice(['vp', 've', 'linear']))
@click.option('--scaling',                 help='Ablate signal scaling s(t)', metavar='vp|none',                    type=click.Choice(['vp', 'none']))

## DG configuration
@click.option('--dg_weight_1st_order',     help='Weight of DG for 1st prediction',       metavar='FLOAT',           type=click.FloatRange(min=0), default=2., show_default=True)
@click.option('--dg_weight_2nd_order',     help='Weight of DG for 2nd prediction',       metavar='FLOAT',           type=click.FloatRange(min=0), default=0., show_default=True)
@click.option('--dg_logit_max', 'dg_logit_max', help='Highest prediction accuracy to apply DG', metavar='FLOAT',    type=click.FloatRange(min=0), default=1.0, show_default=True)
@click.option('--S_clip_min', 'S_clip_min', help='Minimum time[0,1] to apply DG', metavar='FLOAT',                  type=click.FloatRange(min=0), default=0., show_default=True)
@click.option('--S_clip_max', 'S_clip_max', help='Maximum time[0,1] to apply DG', metavar='FLOAT',                  type=click.FloatRange(min=0), default='inf', show_default=True)
# @click.option('--boosting',                help='If true, dg scale up low log ratio samples', metavar='INT',       type=click.IntRange(min=0), default=0, show_default=True)

@click.option('--save_each_sample',        help='Save each sample as image file',                                   is_flag=True, default=False, show_default=True)
## Discriminator checkpoint
@click.option('--discriminator', 'discriminator_pkl', help='Path of discriminator',  metavar='STR', type=str,
              default=None, show_default=True)

def main(network_pkl, discriminator_pkl, outdir, subdirs, seeds, class_idx, max_batch_size, save_each_sample, device=torch.device('cuda'), **sampler_kwargs):
    """Generate random images using the techniques described in the paper
    "Elucidating the Design Space of Diffusion-Based Generative Models".

    Examples:

    \b
    # Generate 64 images and save them as out/*.png
    python generate.py --outdir=out --seeds=0-63 --batch=64 \\
        --network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl

    \b
    # Generate 1024 images using 2 GPUs
    torchrun --standalone --nproc_per_node=2 generate.py --outdir=out --seeds=0-999 --batch=64 \\
        --network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl
    """
    dist.init()
    num_batches = ((len(seeds) - 1) // (max_batch_size * dist.get_world_size()) + 1) * dist.get_world_size()
    all_batches = torch.as_tensor(seeds).tensor_split(num_batches)
    rank_batches = all_batches[dist.get_rank():: dist.get_world_size()]

    # Rank 0 goes first.
    if dist.get_rank() != 0:
        torch.distributed.barrier()

    # Load network.
    dist.print0(f'Loading network from "{network_pkl}"...')
    with dnnlib.util.open_url(network_pkl, verbose=(dist.get_rank() == 0)) as f:
        net = pickle.load(f)['ema'].to(device)
    discriminator = None
    if discriminator_pkl is not None:
        dist.print0(f'Loading discriminator from "{discriminator_pkl}"...')
        with dnnlib.util.open_url(discriminator_pkl, verbose=(dist.get_rank() == 0)) as f:
            discriminator = pickle.load(f)['ema'].to(device)

    # Other ranks follow.
    if dist.get_rank() == 0:
        torch.distributed.barrier()
    drop = []
    pred_x0_npz = []
    # Loop over batches.
    dist.print0(f'Generating {len(seeds)} images to "{outdir}"...')
    for batch_seeds in tqdm.tqdm(rank_batches, unit='batch', disable=(dist.get_rank() != 0)):
        torch.distributed.barrier()
        batch_size = len(batch_seeds)
        if batch_size == 0:
            continue

        # Pick latents and labels.
        rnd = StackedRandomGenerator(device, batch_seeds)
        latents = rnd.randn([batch_size, net.img_channels, net.img_resolution, net.img_resolution], device=device)
        class_labels = None
        if net.label_dim:
            class_labels = torch.eye(net.label_dim, device=device)[
                rnd.randint(net.label_dim, size=[batch_size], device=device)]
        if class_idx is not None:
            class_labels[:, :] = 0
            class_labels[:, class_idx] = 1

        # Generate images.
        sampler_kwargs = {key: value for key, value in sampler_kwargs.items() if value is not None}
        sampler_fn = edm_sampler
        images = sampler_fn(net, latents, class_labels, randn_like=rnd.randn_like,
                            discriminator=discriminator, **sampler_kwargs)

        # Save images.
        if save_each_sample:
            images_np = (images * 127.5 + 128).clip(0, 255).to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()
            for seed, image_np, lab in zip(batch_seeds, images_np, class_labels.argmax(dim=1)):
                image_dir = os.path.join(outdir, f'{seed - seed % 1000:06d}') if subdirs else outdir
                os.makedirs(image_dir, exist_ok=True)
                image_path = os.path.join(image_dir, f'{seed:06d}_{lab.item()}.png')
                if image_np.shape[2] == 1:
                    PIL.Image.fromarray(image_np[:, :, 0], 'L').save(image_path)
                else:
                    PIL.Image.fromarray(image_np, 'RGB').save(image_path)
        else:
            image_dir = os.path.join(outdir, f'{seed - seed % 1000:06d}') if subdirs else outdir
            os.makedirs(image_dir, exist_ok=True)
            image_grid = (images * 127.5 + 128).clip(0, 255).to(torch.uint8)
            image_grid = image_grid.reshape(int(np.sqrt(len(seeds))), int(np.sqrt(len(seeds))),
                                            *image_grid.shape[1:]).permute(0, 3, 1, 4, 2)
            image_grid = image_grid.reshape(int(np.sqrt(len(seeds))) * net.img_resolution,
                                            int(np.sqrt(len(seeds))) * net.img_resolution, net.img_channels)
            image_grid = image_grid.cpu().numpy()
            if image_grid.shape[2] == 1:
                PIL.Image.fromarray(image_grid[:, :, 0], 'L').save(
                    os.path.join(image_dir, f'grid-{class_idx}-{batch_seeds[0]:06d}-{batch_seeds[-1]:06d}.png'))
            else:
                PIL.Image.fromarray(image_grid, 'RGB').save(
                    os.path.join(image_dir, f'grid-{class_idx}-{batch_seeds[0]:06d}-{batch_seeds[-1]:06d}.png'))

    # dist.print0(f'Drop mean: {np.mean(drop)}, std: {np.std(drop)}')
    # Done.
    torch.distributed.barrier()
    dist.print0('Done.')

#----------------------------------------------------------------------------

if __name__ == "__main__":
    main()

#----------------------------------------------------------------------------
