import random
import time
from io import BytesIO

import PIL.Image as Image
import numpy as np
import openvino.runtime as ov
import pandas as pd
import streamlit as st
from cellpose import models
from skimage.color import label2rgb

from openvino_utils import ov_inference

cellpose_model = models.Cellpose(
    gpu=False,
    model_type="cyto2",
    net_avg=False,
)

core = ov.Core()
core_model = core.read_model("./converted_unet_cellpose/cyto2.xml")
ov_model = core.compile_model(core_model, "CPU")

gifs = [
    "https://gifdb.com/images/high/waiting-fingers-agent-cooper-83uxmdriug9b3eph.gif",
    "https://gifdb.com/images/high/i-ll-be-waiting-nacho-libre-ij5d7npbjldtnd2j.webp",
    "https://gifdb.com/images/high/waiting-cat-nail-file-fuyageziynzjynxt.webp",
    "https://gifdb.com/images/high/waiting-for-you-squid-game-f52p9bou56mol0wa.webp",
    "https://gifdb.com/images/high/waiting-man-tick-tock-lucgba3e6vq2eecp.webp",
    "https://gifdb.com/images/high/waiting-impatiently-dumbledore-cu8xf3lc3o7pzyxq.webp",
    "https://gifdb.com/images/high/waiting-dancing-mr-bean-zqemgdq7qldp6jl7.webp",
    "https://gifdb.com/images/high/waiting-sad-pablo-narcos-zz7uiyio8n4g1yra.webp",
    "https://gifdb.com/images/high/still-waiting-justin-timberlake-e06wu78iv4c62mmz.webp",
    "https://gifdb.com/images/high/still-waiting-skeleton-chair-4yggsjnib7cs49ig.webp",
    "https://gifdb.com/images/high/waiting-jared-silicon-valley-0jqzysdhn9om30av.webp",
    "https://gifdb.com/images/high/waiting-for-reply-mocha-oo9sso1g90140bxi.webp",
    "https://gifdb.com/images/high/still-waiting-mickey-mouse-w7kp6rhecp3yizsx.webp",
    "https://gifdb.com/images/high/still-waiting-little-rascals-zedphzadck29jgl6.webp",
    "https://gifdb.com/images/high/still-waiting-boo-monsters-inc-zmx0wumbaimraxf8.webp",
]

st.title("OpenVINO :handshake: CellPose")

st.caption('Developed by CellAI - BECLS')
st.caption('Contributors: Gabriele Aldeghi :pig: Filip Krasniqi :bear:')

st.markdown(
    "[![Repo](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/Valitacell-Ltd)   [![Repo](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/grozby)   [![Repo](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/filipkrasniqi)   [![Repo](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/MouseLand/cellpose)   Copyright © 2020 Howard Hughes Medical Institute",
    unsafe_allow_html=True,
)

uploaded_file = st.file_uploader("Choose a file")
container_warning = st.container()

col1, col2, col3 = st.columns(3)
container_table = st.container()
st.caption("OpenVINO vs CellPose: Inference Time [s] vs Image Size [pixels]")
st.image("assets/cellpose_benchmark.png")

MAX_SIZE = 1024

if uploaded_file is not None:
    try:
        bytes_data = uploaded_file.getvalue()
        image = Image.open(BytesIO(bytes_data))

        img_np = np.asarray(image)
        display_warning = (img_np.shape[0] > MAX_SIZE or
                           img_np.shape[1] > MAX_SIZE)

        if display_warning:
            former_shape = img_np.shape
            img_np = img_np[:MAX_SIZE, :MAX_SIZE]
            container_warning.write(f"WARNING: Image has been cropped "
                                    f"from {former_shape} to {img_np.shape}")

        img_input_cellpose = img_np.copy()

        if len(img_input_cellpose.shape) <= 2:
            img_input_cellpose = np.expand_dims(img_input_cellpose, axis=-1)
        img_input_cellpose = np.expand_dims(img_input_cellpose, axis=0)

        img_input_cellpose = img_input_cellpose / (2**16 - 1)

        col1.write("Input")
        col2.write("CellPose")
        col3.write("OpenVINO")

        col1.image(img_input_cellpose)

        cellpose_img_container = col2.empty()
        cellpose_img_container.image(random.choice(gifs))

        t1 = time.time()

        cp_mask, *_ = cellpose_model.eval(
            img_input_cellpose,
            batch_size=64,
            normalize=False,
            diameter=None,
            flow_threshold=0.4,
            channels=(0, 0),
        )

        t2 = time.time()

        cp_overlay = (label2rgb(
            cp_mask,
            image=img_input_cellpose.squeeze(),
            bg_label=0,
        ) * 255).astype(np.uint8,)

        cellpose_img_container.image(cp_overlay)

        ov_img_container = col3.empty()
        ov_img_container.image(random.choice(gifs))

        img_input_ov = img_np.copy()

        image = np.expand_dims(img_input_ov, axis=0)

        image = np.concatenate(
            [
                image,
                np.zeros_like(image),
            ],
            axis=0,
        ).astype(float)

        image /= 2**16 - 1

        t3 = time.time()
        ov_mask = ov_inference(model=ov_model, x=image)

        t4 = time.time()

        ov_overlay = (label2rgb(
            ov_mask,
            image=img_input_cellpose.squeeze(),
            bg_label=0,
        ) * 255).astype(np.uint8,)

        ov_img_container.image(ov_overlay)

        df = pd.DataFrame([
            {
                "Model": "CellPose",
                "Execution Time [s]": f"{(t2-t1):.2f} seconds"
            },
            {
                "Model": "OpenVINO",
                "Execution Time [s]": f"{(t4-t3):.2f} seconds"
            },
        ])

        container_table.table(df)
    except Exception as e:
        container_warning.write("WARNING: an error occurred. Please retry.")