import streamlit as st
from pathlib import Path
from streamlit import session_state as sess
import base64

from utils import raw_image, overlay_mask_on_image, image_to_file_path, AreaModel

import streamlit.components.v1 as components

st.set_page_config(layout="wide") 
# Set the title of the app
st.write("### Cumulus Oocyte Complex (COC) area estimation")
# Add a description
st.write("This demo allows the computation of the area in pixels of a cumulus oocyte complex from its photograph implementing the method described in ")
st.write("""Athanasiou, Georgios, Jesus Cerquides, Annelies Raes, Nima Azari-Dolatabad, Daniel Angel-Velez, Ann Van Soom, and Josep-Lluis Arcos. <a href="https://doi.org/10.3233/FAIA220346">Detecting the Area of Bovine Cumulus Oocyte Complexes Using Deep Learning and Semantic Segmentation</a>. In Artificial Intelligence Research and Development, 249–58. IOS Press, 2022.""", unsafe_allow_html=True)
st.write("Upload the photograph of your oocyte here:")
if 'image_path' not in sess:
    sess.image_path = None
if 'last_image_path' not in sess:
    sess.last_image_path = None
if 'model' not in sess:
    sess.model = AreaModel()
if 'image' not in sess:
    sess.image = None

def compute_area():
    if sess.last_image_path != sess.image_path:
        with (st.spinner("In progress, should not be more than 10 seconds...")):
            sess.img = raw_image(sess.image_path)
            sess.area, sess.roi_img, sess.mask = sess.model.compute_area(sess.img)
            sess.last_image_path = sess.image_path
            return

def update_from_uploader():
    uploaded_file = sess["uploaded_file"]
    if uploaded_file is not None:
        sess.image_path = image_to_file_path(uploaded_file)
        compute_area()

image = st.file_uploader("Oocyte image", type=["jpg", "jpeg", "png"] , key="uploaded_file", on_change=update_from_uploader)

if sess.image_path is not None:
    st.write(f"### Area computed = {sess.area:.2f} pixels")
    col1, col2, col3 = st.columns(3)
    with col1:
        st.image(sess.img, caption='Original oocyte', use_column_width=True)
    with col2:
        st.image(sess.roi_img, caption='ROI', use_column_width=True)
    with col3:
        st.image(overlay_mask_on_image(sess.roi_img, sess.mask), caption='ROI with the deep learning computed mask', use_column_width=True, clamp=True)

# Show examples

st.write("Or select one of the examples below:")
examples = ["app/static/oocytes/immature/1.1.png",
            "app/static/oocytes/immature/14.10.png",
            "app/static/oocytes/immature/13.4.png"]
examples_to_load = [str(Path(*Path(ex).parts[1:])) for ex in examples]
example_imgs = [raw_image(ex) for ex in examples_to_load]
cols = st.columns(len(examples))
for i, col in enumerate(cols):
    with col:
        def set_example(i):
            def f():
                sess.image_path = examples_to_load[i]
                # print(sess.image_path)
                compute_area()
            return f
        st.button("Example "+str(i+1), on_click=set_example(i))
        st.image(example_imgs[i], caption='Example ' + str(i+1), use_column_width=True)

footer_html = """<div style='text-align: center;'>
  <p><img src="./app/static/img/EurovaLogo.png" height=100 alt="EUROVA logo"/>
  <p>Developed with ❤️ at <a href="https://iiia.csic.es"><img src="./app/static/img/iiia-logo.png" height=30 alt="IIIA logo"/> IIIA - CSIC</a> by <a href="https://www.iiia.csic.es/~cerquide">Jesus Cerquides</a>, Giorgios Athanasiou and Josep LLuís Arcos </p>
  <p>In collaboration with <a href="https://www.ugent.be/en"><img src="./app/static/img/Ghent.png" height=40 alt="UGhent logo"/> Ghent University</a></p>
  <p><img src="./app/static/img/MSCA.jpg" height=30 alt="MSCA logo"/> Funded by the European Union Horizon 2020 research and innovation programme under the Marie Sklodowka-Curie grant agreement 860960
</div>"""
st.markdown(footer_html, unsafe_allow_html=True)