from google.colab import auth

auth.authenticate_user()

project_id = "probml"
!gcloud config set project {project_id}

"""
this should be the format of the checkpoint filetree:
  checkpoint_path >> model(optimizer)_checkpoint_file.
  checkpoint_path_ema >> ema_checkpoint_file
"""
checkpoint_path = "/content/vdvae_cifar10_2.86/latest_cifar10"

# checkpoints are downloaded at these paths.
#  vdvae_cifar10_2.86/latest_cifar10 - optimizer+mode
#  vdvae_cifar10_2.86/latest_cifar10_ema - ema_params'

# @title Download checkpoints
!gsutil cp -r gs://gsoc_bucket/vdvae_cifar10_2.86 ./
!ls -l /content/vdvae_cifar10_2.86/latest_cifar10
!ls -l /content/vdvae_cifar10_2.86/latest_cifar10_ema

!git clone https://github.com/j-towns/vdvae-jax.git

%cd vdvae-jax

!pip install --quiet flax

import os

try:
    os.environ["COLAB_TPU_ADDR"]
    import jax.tools.colab_tpu

    jax.tools.colab_tpu.setup_tpu()
except:
    pass

import jax

jax.local_devices()

"""
Explanation: <a href="https://colab.research.google.com/github/always-newbie161/probml-notebooks/blob/jax_vdvae/notebooks/vdvae_jax_cifar_demo.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
This notebook shows demo working with vdvae in jax and the code used is from vdvae-jax from Jamie Townsend
Setup
End of explanation
"""


from hps import HPARAMS_REGISTRY, Hyperparams, add_vae_arguments
from train_helpers import setup_save_dirs
import argparse
import dataclasses

H = Hyperparams()
parser = argparse.ArgumentParser()
parser = add_vae_arguments(parser)
parser.set_defaults(hps="cifar10", conv_precision="highest")

H = dataclasses.replace(H, **vars(parser.parse_args([])))
hparam_sets = [x for x in H.hps.split(",") if x]
for hp_set in hparam_sets:
    hps = HPARAMS_REGISTRY[hp_set]
    parser.set_defaults(**hps)

H = dataclasses.replace(H, **vars(parser.parse_args([])))
H = setup_save_dirs(H)

"""
Explanation: Model
(for cifar10)
Setting up hyperparams
End of explanation
"""


hparams = dataclasses.asdict(H)
for k in ["enc_blocks", "dec_blocks", "zdim", "n_batch", "device_count"]:
    print(f"{k}:{hparams[k]}")

from utils import logger
from jax.interpreters.xla import DeviceArray

log = logger(H.logdir)
if H.log_wandb:
    import wandb

    def logprint(*args, pprint=False, **kwargs):
        if len(args) > 0:
            log(*args)
        wandb.log({k: np.array(x) if type(x) is DeviceArray else x for k, x in kwargs.items()})

    wandb.init(config=dataclasses.asdict(H))
else:
    logprint = log

import numpy as np
from jax import lax
import torch
import imageio
from PIL import Image
import glob
from torch.utils.data import DataLoader
from torchvision import transforms


np.random.seed(H.seed)
torch.manual_seed(H.seed)
H = dataclasses.replace(
    H,
    conv_precision={"default": lax.Precision.DEFAULT, "high": lax.Precision.HIGH, "highest": lax.Precision.HIGHEST}[
        H.conv_precision
    ],
    seed_init=H.seed,
    seed_sample=H.seed + 1,
    seed_train=H.seed + 2 + H.host_id,
    seed_eval=H.seed + 2 + H.host_count + H.host_id,
)
print("training model on ", H.dataset)

"""
Explanation: This model is a hierarchical model with multiple stochastic blocks with multiple deterministic layers. You can know about model skeleton by observing the encoder and decoder "strings"
How to understand the string:
*   blocks are comma seperated 
*   axb implies there are b res blocks(set of Conv layers) for dimensions axa
*  amb implies it is a mixin block which increases the inter-image dims from a to b using nearest neighbour upsampling (used in decoder)
* adb implies it's a block with avg-pooling layer which reduces the dims from a to b(used in encoder)
for more understanding refer to this paper
End of explanation
"""


!./setup_cifar10.sh

"""
Explanation: Downloading cifar10 dataset
End of explanation
"""


from data import set_up_data

H, data_train, data_valid_or_test, preprocess_fn = set_up_data(H)

from train_helpers import load_vaes

H = dataclasses.replace(H, restore_path=checkpoint_path)
optimizer, ema_params, start_epoch = load_vaes(H, logprint)

start_epoch  # no.of.epochs trained

# Hparams for the current model
hparams = dataclasses.asdict(H)
for i, k in enumerate(sorted(hparams)):
    logprint(f"type=hparam, key={k}, value={getattr(H, k)}")

"""
Explanation: Setting up the model, data and the preprocess fn.
End of explanation
"""


from train import run_test_eval

run_test_eval(H, ema_params, data_valid_or_test, preprocess_fn, logprint)

"""
Explanation: Evaluation
End of explanation
"""


def zoom_in(fname, shape):
    im = Image.open(fname)
    resized_im = im.resize(shape)
    resized_im.save(fname)


def save_n_show(images, order, image_shape, fname, zoom=True, show=False):
    n_rows, n_images = order
    im = (
        images.reshape((n_rows, n_images, *image_shape))
        .transpose([0, 2, 1, 3, 4])
        .reshape([n_rows * image_shape[0], n_images * image_shape[1], 3])
    )
    print(f"printing samples to {fname}")
    imageio.imwrite(fname, im)

    if zoom:
        zoom_in(fname, (640, 64))  # w=640, h=64

    if show:
        display(Image.open(fname))

"""
Explanation: Function to save and show of batch of images given as a numpy array.
End of explanation
"""


n_images = 10
num_temperatures = 3
image_shape = [H.image_size, H.image_size, H.image_channels]
H = dataclasses.replace(H, num_images_visualize=n_images, num_temperatures_visualize=num_temperatures)

"""
Explanation: Generations
End of explanation
"""


H.save_dir

"""
Explanation: Images will be saved in the following dir
End of explanation
"""


from jax import random
from vae import VAE
from flax import jax_utils
from functools import partial

rng = random.PRNGKey(H.seed_sample)

ema_apply = partial(VAE(H).apply, {"params": jax_utils.unreplicate(ema_params)})

forward_uncond_samples = partial(ema_apply, method=VAE(H).forward_uncond_samples)

temperatures = [1.0, 0.9, 0.8, 0.7]
for t in temperatures[: H.num_temperatures_visualize]:
    im = forward_uncond_samples(n_images, rng, t=t)
    im = np.asarray(im)
    save_n_show(im, [1, n_images], image_shape, f"{H.save_dir}/generations-tem-{t}.png")

for t in temperatures[: H.num_temperatures_visualize]:
    print("=" * 25)
    print(f"Generation of {n_images} new images for t={t}")
    print("=" * 25)
    fname = f"{H.save_dir}/generations-tem-{t}.png"
    display(Image.open(fname))

"""
Explanation: As the model params are replicated over multiple devices, unreplicated copy of them is made to use it for sampling and generations.
End of explanation
"""


n_images = 10
image_shape = [H.image_size, H.image_size, H.image_channels]

"""
Explanation: Reconstructions
End of explanation
"""


from train import get_sample_for_visualization

viz_batch_original, viz_batch_processed = get_sample_for_visualization(
    data_valid_or_test, preprocess_fn, n_images, H.dataset
)

"""
Explanation: Preprocessing images before getting the latents
End of explanation
"""


forward_get_latents = partial(ema_apply, method=VAE(H).forward_get_latents)
forward_samples_set_latents = partial(ema_apply, method=VAE(H).forward_samples_set_latents)

"""
Explanation: Getting the partial functions from the model methods
End of explanation
"""


zs = [s["z"] for s in forward_get_latents(viz_batch_processed, rng)]

"""
Explanation: Getting latents of different levels.
End of explanation
"""


recons = []
lv_points = np.floor(np.linspace(0, 1, H.num_variables_visualize + 2) * len(zs)).astype(int)[1:-1]
for i in lv_points:
    recons.append(forward_samples_set_latents(n_images, zs[:i], rng, t=0.1))

"""
Explanation: No of latent observations used depends on H.num_variables_visualize, altering it gives different resolutions of the reconstructions.
End of explanation
"""


orig_im = np.array(viz_batch_original)
print("Original test images")
save_n_show(orig_im, [1, n_images], image_shape, f"{H.save_dir}/orig_test.png", show=True)

"""
Explanation: Original Images
End of explanation
"""


for i, r in enumerate(recons):
    r = np.array(r)
    print("=" * 25)
    print(f"Generation of {n_images} new images for {i+1}x resolution")
    print("=" * 25)
    fname = f"{H.save_dir}/recon_test-res-{i+1}x.png"
    save_n_show(r, [1, n_images], image_shape, fname, show=True)

"""
Explanation: Reconstructions.
End of explanation
"""


