import sys, os
import gradio as gr
import numpy as np
sys.path.append('stylegan3')
import utils

def to_uint8(im, ndim=2):
    im -= np.min(im)
    im /= np.max(im)
    im *= 255.
    im = np.asarray(im, dtype=np.uint8)
    if ndim == 3:    
        if im.ndim == 2:
            im = np.expand_dims(im, axis=-1)
        elif im.ndim == 3:
            if im.shape[0] == 1: 
                np.transpose(im, (1,2,0))
        im = np.tile(im, (1,1,3)) #make fake RGB
        return im
    elif ndim ==2:
        if im.ndim == 2:
            return im
        if im.ndim == 3:
            if im.shape[0] == 1:  #[1, H, W]
                return im[0,...]
            elif im.shape[2] == 1: #[H, W, 1]
                return im[...,0]
            else:
                raise AssertionError(f"Unexpected image passed to to_uint8 with shape: {np.shape(im)}.")


in_gpu = False
num_images = 1
G = utils.load_default_gen(in_gpu=in_gpu)
sampler = utils.SampleFromGAN(G=G, z_shp=[num_images, G.z_dim], in_gpu=in_gpu)

def sample_GAN():
    im = sampler()
    im = im.numpy()
    im = np.transpose(im, (1,2,0))
    im = np.squeeze(im) #if single channel (yes), drop it.
    # print(f"sample_linearBP: im shape: {im.shape}; min: {np.min(im)}, max: {np.max(im)}.")
    im = to_uint8(im, ndim=2)
    # print(f'1. uint image shape: {im.shape}')
    return im


title="Generate fake linear array images"
description="Generate fake linear array images."

with gr.Blocks() as demo:
    gr.Markdown(description)
    # with gr.Row(equal_height=True):
    #     button_gen = gr.Button("Generate fake linear image")
    # with gr.Row(equal_height=True):
    #     output_im = gr.Image(type="numpy", shape=(256, 256), image_mode="L", label="fake image", interactive=False) #grayscale image
    with gr.Column(scale=1):
        button_gen = gr.Button("Generate fake linear image")
    with gr.Column(scale=2):
        output_im = gr.Image(type="numpy", shape=(256, 256), image_mode="L", label="fake image", interactive=False) #grayscale image
    button_gen.click(sample_GAN, inputs=None, outputs=output_im)
    
demo.launch(share=False, show_tips=True, enable_queue=True)