''' ----------------------------------------
* 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_allRnd_0317()
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;">ML-SIM: Reconstruction of SIM images with deep learning</h1>'

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

### <a href="https://opg.optica.org/boe/viewmedia.cfm?uri=boe-12-5-2720&html=true" target='_blank' > ML-SIM: universal reconstruction of structured illumination microscopy images using transfer learning </a>

_Charles N. Christensen<sup>1,2,*</sup>, Edward N. Ward<sup>1</sup>, Meng Lu<sup>1</sup>, Pietro Lio<sup>2</sup>, Clemens F. Kaminski_</br></br>
<sup>1</sup>University of Cambridge, Department of Chemical Engineering and Biotechnology, Laser Analytics Group</br>
<sup>2</sup>University of Cambridge, Department of Computer Science and Technology, Artificial Intelligence Group</br>
- GitHub: [charlesnchr](http://github.com/charlesnchr)
- Email: charles.n.chr@gmail.com
- Publication: <a href='https://opg.optica.org/boe/viewmedia.cfm?uri=boe-12-5-2720&html=true'
target='_blank'>Journal</a>, <a href='https://arxiv.org/abs/2003.11064' target='_blank'>Preprint</a>
---

## 🔬 To run ML-SIM
Upload a TIFF image and hit submit or select one from the examples below. Note that the model here is trained for SIM stacks of 9 frames (3x3 configuration, i.e. 3 phase steps for 3 orientations).
"""

article = """
---
Select an example from the list above to try the model on a test image. Below you can see what the file names correspond to.

The ML-SIM test images are from two different microscopes (MAI-SIM and SLM-SIM) in addition to two simulated images.

Wide-field projections in pseudo-colour of these examples are shown below:

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

---
### Read more
- <a href='https://ML-SIM.com' target='_blank'>ML-SIM.com</a>
- <a href='https://charles-christensen.com' target='_blank'>Website</a>
- <a href='https://github.com/charlesnchr/ML-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 (ML-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()