import gradio as gr
import numpy as np
import time


def predict(x):
    return np.fliplr(x)


def compress():
    time.sleep(1)
    return 'The model has been compressed successfully.'


with gr.Blocks() as demo:
    with gr.Column():
        gr.Radio(["Image classification", "Object detection", "Semantic segmentation"], label="Tasks"),
        gr.Radio(["ResNet", "VGG", "MobileNet"], label="Models"),
        gr.Radio(["Weight quantization","Knowledge distillation","Network pruning", "Neural Architecture Search"],
                 label="Compression methods"),
        gr.Radio(["Jetson Nano"], label="Deployments")
        compress_btn = gr.Button("compress")
        output_compress = gr.Textbox(lines=1, label="Model Compression Results", visible=False)
        
    with gr.Row():
        Original_config = gr.Dataframe(headers=["#Params.(M)", "FLOPs(G)"], datatype=[
                     "str", "str"], row_count=1, value=[['63.8M','250G']], label="Original model config", visible=False)
        Compressed_config = gr.Dataframe(headers=["#Params.(M)", "FLOPs(G)"], datatype=[
                     "str", "str"], row_count=1,  value=[['34.6M','126G']],label="Compressed model config", visible=False)
    with gr.Row():
        input_predict = gr.Image(label="input")
        output_predict = gr.Image(label="output")
    predict_btn = gr.Button("predict")
    state = gr.State()
    compress_btn.click(fn=compress, inputs=None,
                       outputs=output_compress, api_name="compress")
    compress_btn.click(lambda : (gr.Textbox.update(visible=True), "visible"), None, [output_compress, state])
    output_compress.change(lambda: (Original_config.update(visible=True), "visible"), None, [Original_config, state])   
    output_compress.change(lambda: (Compressed_config.update(visible=True), "visible"), None, [Compressed_config, state])   
    predict_btn.click(fn=predict, inputs=input_predict,
                      outputs=output_predict, api_name="predict")
    
demo.launch(share=True)