import math

import fire
import jax
import jax.numpy as jnp
import numpy as np
import opax
import pax
import tensorflow as tf
from PIL import Image

from data_loader import load_celeb_a
from model import GaussianDiffusion, UNet


def make_image_grid(images, padding=2):
    """Place images in a square grid."""
    n = images.shape[0]
    size = int(math.sqrt(n))
    assert size * size == n, "expecting a square grid"
    img = images[0]

    H = img.shape[0] * size + padding * (size + 1)
    W = img.shape[1] * size + padding * (size + 1)
    out = np.zeros((H, W, img.shape[-1]), dtype=img.dtype)
    for i in range(n):
        x = i % size
        y = i // size
        xstart = x * (img.shape[0] + padding) + padding
        xend = xstart + img.shape[0]
        ystart = y * (img.shape[1] + padding) + padding
        yend = ystart + img.shape[1]
        out[xstart:xend, ystart:yend, :] = images[i]
    return out


def train(
    batch_size: int = 32,
    learning_rate: float = 1e-4,
    num_training_steps: int = 10_000,
    log_freq: int = 1000,
    image_size: int = 64,
    random_seed: int = 42,
):

    pax.seed_rng_key(random_seed)

    model = UNet(dim=64, dim_mults=(1, 2, 4, 8))

    diffusion = GaussianDiffusion(
        model,
        image_size=image_size,
        timesteps=1000,
        loss_type="l1",  # L1 or L2
    )

    dataset = load_celeb_a()

    dataloader = (
        dataset
        # .cache()  # TODO: memory leak?
        .repeat()
        .shuffle(batch_size * 100)
        .batch(batch_size)
        .take(num_training_steps)
        .prefetch(tf.data.AUTOTUNE)
    )

    def loss_fn(model, inputs):
        model, loss = pax.module_and_value(model)(inputs)
        return loss, (loss, model)

    update_fn = pax.utils.build_update_fn(loss_fn)
    fast_update_fn = jax.jit(update_fn)

    optimizer = opax.adam(learning_rate)(diffusion.parameters())

    from tqdm.auto import tqdm

    total_loss = 0.0
    tr = tqdm(dataloader)
    for step, batch in enumerate(tr, 1):
        batch = jax.tree_map(lambda x: x.numpy(), batch)
        diffusion, optimizer, loss = fast_update_fn(diffusion, optimizer, batch)
        total_loss = total_loss + loss

        if step % log_freq == 0:
            loss = total_loss / log_freq
            total_loss = 0.0
            tr.write(f"[step {step:05d}]  train loss {loss:.3f}")

            imgs = jax.device_get(diffusion.eval().sample(16))
            imgs = ((imgs * 0.5 + 0.5) * 255).astype(jnp.uint8)
            imgs = make_image_grid(imgs)
            im = Image.fromarray(imgs)
            im.save(f"sample_{step:05d}.png")


if __name__ == "__main__":
    fire.Fire(train)
