import streamlit as st

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import warnings
warnings.filterwarnings('ignore')
# %matplotlib inline

import tensorflow
print (tensorflow.__version__)

st.header("Welcome to the Generative Playground")

from tensorflow.keras.datasets import mnist,cifar10

option = st.selectbox(
    "Which model would you like to get prediction with?",
    ("None","Auto-Regressor", "Auto-Encoder", "Diffusion-Model","Other"))

st.write("You selected:", option)

if option == "None":
    st.write("Please Select the model to get the fun prediction.... :)")

if option == "Auto-Encoder":
    st.write("It is under development")
    st.write("Stay tune... Comming soon... :)")
    
if option == "Other":
    st.write("Stay tune... Updating soon... :)")

if option == "Diffusion-Model":
    st.write("It is under development")
    st.write("Stay tune... Comming soon... :)")
    
if option == "Auto-Regressor":
    if st.button("Run"):
        st.write("Running Auto-Regressor")

        st.write("trained on --> cifar-10 dataset, RTX-GPU's, 50-epochs")
        st.write("This is trail model, updated version will be updated consicutively.")
        
        (trainX, trainy), (testX, testy) = cifar10.load_data()
        
        print('Training data shapes: X=%s, y=%s' % (trainX.shape, trainy.shape))
        print('Testing data shapes: X=%s, y=%s' % (testX.shape, testy.shape))
        
        
        
        for k in range(4):
            fig = plt.figure(figsize=(9,6))
            for j in range(9):
                i = np.random.randint(0, 10000)
                plt.subplot(990 + 1 + j)
                plt.imshow(trainX[i], cmap='gray_r')
                # st.pyplot(fig)
                plt.axis('off')
                #plt.title(trainy[i])
            plt.show()
            st.pyplot(fig)
        
        
        # asdfaf
        
        trainX = np.where(trainX < (0.33 * 256), 0, 1)
        train_data = trainX.astype(np.float32)
        
        testX = np.where(testX < (0.33 * 256), 0, 1)
        test_data = testX.astype(np.float32)
        
        train_data = np.reshape(train_data, (50000, 32, 32, 3))
        test_data = np.reshape(test_data, (10000, 32, 32, 3))
        
        print (train_data.shape, test_data.shape)
        
        
        import tensorflow
        
        class PixelConvLayer(tensorflow.keras.layers.Layer):
            def __init__(self, mask_type, **kwargs):
                super(PixelConvLayer, self).__init__()
                self.mask_type = mask_type
                self.conv = tensorflow.keras.layers.Conv2D(**kwargs)
        
            def build(self, input_shape):
                # Build the conv2d layer to initialize kernel variables
                self.conv.build(input_shape)
                # Use the initialized kernel to create the mask
                kernel_shape = self.conv.kernel.get_shape()
                self.mask = np.zeros(shape=kernel_shape)
                self.mask[: kernel_shape[0] // 2, ...] = 1.0
                self.mask[kernel_shape[0] // 2, : kernel_shape[1] // 2, ...] = 1.0
                if self.mask_type == "B":
                    self.mask[kernel_shape[0] // 2, kernel_shape[1] // 2, ...] = 1.0
        
            def call(self, inputs):
                self.conv.kernel.assign(self.conv.kernel * self.mask)
                return self.conv(inputs)
        
        
        # Next, we build our residual block layer.
        # This is just a normal residual block, but based on the PixelConvLayer.
        class ResidualBlock(tensorflow.keras.layers.Layer):
            def __init__(self, filters, **kwargs):
                super(ResidualBlock, self).__init__(**kwargs)
                self.conv1 = tensorflow.keras.layers.Conv2D(
                    filters=filters, kernel_size=1, activation="relu"
                )
                self.pixel_conv = PixelConvLayer(
                    mask_type="B",
                    filters=filters // 2,
                    kernel_size=3,
                    activation="relu",
                    padding="same",
                )
                self.conv2 = tensorflow.keras.layers.Conv2D(
                    filters=filters, kernel_size=1, activation="relu"
                )
        
            def call(self, inputs):
                x = self.conv1(inputs)
                x = self.pixel_conv(x)
                x = self.conv2(x)
                return tensorflow.keras.layers.add([inputs, x])
        
        inputs = tensorflow.keras.Input(shape=(32,32,3))
        x = PixelConvLayer(
            mask_type="A", filters=128, kernel_size=7, activation="relu", padding="same"
        )(inputs)
        
        for _ in range(5):
            x = ResidualBlock(filters=128)(x)
        
        for _ in range(2):
            x = PixelConvLayer(
                mask_type="B",
                filters=128,
                kernel_size=1,
                strides=1,
                activation="relu",
                padding="valid",
            )(x)
        
        out = tensorflow.keras.layers.Conv2D(
            filters=3, kernel_size=1, strides=1, activation="sigmoid", padding="valid"
        )(x)
        
        pixel_cnn = tensorflow.keras.Model(inputs, out)
        pixel_cnn.summary()
        
        adam = tensorflow.keras.optimizers.Adam(learning_rate=0.0005)
        pixel_cnn.compile(optimizer=adam, loss="binary_crossentropy")
        
        
        # %%
        import os
        checkpoint_path = "training_1/cp.ckpt"
        # checkpoint_path = "training_1/cp.weights.h5"
        checkpoint_dir = os.path.dirname(checkpoint_path)
        
        
        pixel_cnn.load_weights(checkpoint_path)
        
        
        # %% [markdown]
        # # Display Results 81 images
        
        # %%
        # from IPython.display import Image, display
        from tqdm import tqdm
        
        
        # Create an empty array of pixels.
        batch = 1
        pixels = np.zeros(shape=(batch,) + (pixel_cnn.input_shape)[1:])
        batch, rows, cols, channels = pixels.shape
        
        print(pixels.shape)
        
        
        import time 
        
        # progress_text = "Operation in progress. Please wait."
        # my_bar = st.progress(0, progress_text)
        st.caption("Generating..... pls.. wait.. :)")
        my_bar = st.progress(0)
        
        
        # Iterate over the pixels because generation has to be done sequentially pixel by pixel.
        for row in tqdm(range(rows)):
            for col in range(cols):
                for channel in range(channels):
                    time.sleep(0.01)
                    # Feed the whole array and retrieving the pixel value probabilities for the next
                    # pixel.
                    probs = pixel_cnn.predict(pixels)[:, row, col, channel]
                    # Use the probabilities to pick pixel values and append the values to the image
                    # frame.
                    pixels[:, row, col, channel] = tensorflow.math.ceil(
                      probs - tensorflow.random.uniform(probs.shape)
                    )
            my_bar.progress(int(row*3.125))
            # if row<rows/2:
            #     my_bar.progress((rows+1)*2)
            # else:
            #     my_bar.progress(row+51)
            
        my_bar.progress(100)
        time.sleep(1)

        
        from PIL import Image
        # figout = plt.figure(figsize=(9,6))
        # st.image(Image.fromarray((pixels[-1] * 255).astype(np.uint8), 'RGB').show(),caption="image")
        # Convert the generated pixel array to an image
        generated_image = Image.fromarray((pixels[-1] * 255).astype(np.uint8), 'RGB')
        
        # Display the image using Streamlit
        st.image(generated_image, caption="Generated Image")
        
        # counter = 0
        # for i in range(4):
        #     figout = plt.figure(figsize=(9,6))
        #     for j in range(9):
        #         plt.subplot(990 + 1 + j)
        #         plt.imshow(pixels[counter,:,:,0])#, cmap='gray_r')
        #         counter += 1
        #         plt.axis('off')
        #     plt.show()
        #     st.pyplot(figout)
        
        # %%
# else:
#     st.write("Not Available")