''' ----------------------------------------
* Creation Time : Sun Aug 28 21:38:58 2022
* Last Modified : Sun Aug 28 21:41:36 2022
* Author : Charles N. Christensen
* Github : github.com/charlesnchr
----------------------------------------'''

from turtle import title
import gradio as gr
import numpy as np
from PIL import Image
import io
import base64
import skimage
from NNfunctions import *

opt = GetOptions_Swin_2702()
net = LoadModel(opt)

gr.close_all()

def predict(imagefile):
    # img = np.array(skimage.io.imread(imagefile.name))
    # img = np.concatenate((img,img,img),axis=2)
    # img = np.transpose(img, (2,0,1))

    img = skimage.io.imread(imagefile.name)

    # sr,wf,out = EvaluateModel(net,opt,img,outfile)
    sr, wf, sr_download = EvaluateModel(net,opt,img)

    return wf, sr, sr_download

def process_example(filename):
    basename = os.path.basename(filename)
    basename = basename.replace('.png','.tif')
    img = skimage.io.imread('TestImages/%s' % basename)

    sr, wf, sr_download = EvaluateModel(net,opt,img)

    return wf, sr

title = '<h1 style="text-align: center;">VSR-SIM: Spatio-temporal reconstruction method for SIM using vision transformer</h1>'

description = """
This space demonstrates the VSR-SIM method for reconstruction of structured illumination microscopy images.

_Charles N. Christensen<sup>1,2
- GitHub: [charlesnchr](http://github.com/charlesnchr)
- Email: charles.n.chr@gmail.com
- Publication: <a href='https://arxiv.org/abs/2203.00030' target='_blank'>Preprint</a>
---

## 🔬 To run VSR-SIM
Upload a TIFF image and hit submit or select one from the examples below.
"""

article = """

![Example test images](https://i.imgur.com/Lidkbib.jpeg "Example test image for VSR-SIM")

---
### Read more
- <a href='https://VSR-SIM.github.io' target='_blank'>VSR-SIM.github.io</a>
- <a href='https://charles-christensen.com' target='_blank'>Website</a>
- <a href='https://github.com/charlesnchr/VSR-SIM' target='_blank'>Github</a>
- <a href='https://twitter.com/charlesnchr' target='_blank'>Twitter</a>
"""

# inputs = gr.inputs.Image(label="Upload a TIFF image", type = 'pil', optional=False)

inputs = gr.inputs.File(label="Upload a TIFF image", type = 'file', optional=False)
outputs = [
    gr.outputs.Image(label="INPUT (Wide-field projection)"),
    gr.outputs.Image(label="OUTPUT (VSR-SIM)"),
    gr.outputs.File(label="Download SR image" )
    # , gr.outputs.Textbox(type="auto",label="Pet Prediction")
]

examples = glob.glob('*.tif')

interface = gr.Interface(fn=predict,
    inputs=inputs,
    outputs=outputs,
    title = title,
    description=description,
    article=article,
    examples=examples,
    allow_flagging='never',
    cache_examples=False
    )
interface.launch()


# with gr.Blocks() as interface:
#     gr.Markdown(title)
#     gr.Markdown(description)

#     with gr.Row():
#         input1 = gr.inputs.File(label="Upload a TIFF image", type = 'file', optional=False)

#     submit_btn = gr.Button("Reconstruct")

#     with gr.Row():
#         output1 = gr.outputs.Image(label="Wide-field projection")
#         output2 = gr.outputs.Image(label="SIM Reconstruction")

#     output3 = gr.File(label="Download SR image", visible=False)

#     submit_btn.click(
#         predict,
#         input1,
#         [output1, output2, output3]
#     )

#     gr.Examples(examples, input1, [output1, output2, output3])


# interface.launch()