import streamlit as st
from PIL import Image
import jax
import jax.numpy as jnp  # JAX NumPy
import numpy as np
from huggingface_hub import HfFileSystem
from flax.serialization import msgpack_restore, from_state_dict
import time
from generator import Generator, LATENT_DIM
import math

generator = Generator()
variables = generator.init(jax.random.PRNGKey(0), jnp.zeros([1, LATENT_DIM]), training=False)

fs = HfFileSystem()
with fs.open("PrakhAI/AIPlane2/g_checkpoint.msgpack", "rb") as f:
  g_state = from_state_dict(variables, msgpack_restore(f.read()))

def sample_latent(batch, key):
  return jax.random.normal(key, shape=(batch, LATENT_DIM))

def to_img(normalized):
  return ((normalized+1)*255./2.).astype(np.uint8)

st.write("The model and its details are at https://huggingface.co/PrakhAI/AIPlane2")
if st.button('Generate Random'):
  st.session_state['generate'] = None

ROWS = 4
COLUMNS = 4

def set_latent(latent):
  st.session_state['generate'] = latent

if 'generate' in st.session_state:
  unique_id = int(1_000_000 * time.time())
  latents = sample_latent(ROWS * COLUMNS, jax.random.PRNGKey(unique_id))
  previous = st.session_state['generate']
  if previous is not None:
    if "similarity" not in st.session_state:
      st.session_state["similarity"] = 0.5
    similarity = st.number_input(label="Mutation (for \"Generate Similar\") - lower value generates more similar images", key="similarity", min_value=0.01, max_value=1.0)
    latents = np.repeat([previous], repeats=16, axis=0) + similarity * latents
  (g_out128, _, _, _, _, _) = generator.apply({'params': g_state['params'], 'batch_stats': g_state['batch_stats']}, latents, training=False)
  img = np.array(to_img(g_out128))
  for row in range(ROWS):
    with st.container():
      for (col_idx, col) in enumerate(st.columns(COLUMNS)):
        with col:
          idx = row*COLUMNS + col_idx
          st.image(Image.fromarray(img[idx]))
          st.button(label="Generate Similar", key="%d_%d" % (unique_id, idx), on_click=set_latent, args=(latents[idx],))