"""Application to demo inpainting, Median and Bilateral Blur using streamlit.

Run using: streamlit run 10_04_image_restoration_app.py
"""

import streamlit as st
import pathlib
from streamlit_drawable_canvas import st_canvas
import cv2
import numpy as np
import io
import base64
from PIL import Image


# # Function to create a download link for output image
# def get_image_download_link(img, filename, text):
#     """Generates a link to download a particular image file."""
#     buffered = io.BytesIO()
#     img.save(buffered, format='JPEG')
#     img_str = base64.b64encode(buffered.getvalue()).decode()
#     href = f'<a href="data:file/txt;base64,{img_str}" download="{filename}">{text}</a>'
#     return href


# Set title.
st.sidebar.title('Image Restoration App with OpenCV')

# Description
st.sidebar.text('Upload an image and apply various restoration techniques.')


# Specify canvas parameters in application
uploaded_file = st.sidebar.file_uploader("Upload Image to restore:", type=["png", "jpg"])
image = None
res = None

if uploaded_file is not None:

    # Convert the file to an opencv image.
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, 1)
    # Display the uploaded image
    st.subheader("Original Image")
    st.image(image[:, :, ::-1])
    # Display a selection box for choosing the filter to apply.
    option = st.sidebar.selectbox('Median or Bilateral Blur or Inpaint?', ('None', 'Median Blur', 'Bilateral Blur', 'Image Inpaint'))

    if option == 'Median Blur':
        ksize = st.sidebar.slider("ksize: ", 3, 15, 5, 2)
        image = cv2.medianBlur(image, ksize)
        res=image[:,:,::-1]
        st.subheader("Median Blurred Image")
        st.image(res)
        # Display download button
        result = Image.fromarray(res)
        buffered = io.BytesIO()
        result.save(buffered, format="PNG")
        img_bytes = buffered.getvalue()
        st.download_button(label='Download Output', data=img_bytes, file_name='median_blur_output.png', mime='image/png')
        
    elif option == 'Bilateral Blur':
        dia = st.sidebar.slider("diameter: ", 1, 50, 20)
        sigmaColor = st.sidebar.slider("sigmaColor: ", 0, 250, 200, 10)
        sigmaSpace = st.sidebar.slider("sigmaSpace: ", 0, 250, 100, 10)
        image = cv2.bilateralFilter(image, dia, sigmaColor, sigmaSpace)
        res=image[:,:,::-1]
        st.subheader("Bilateral Blurred Image")
        st.image(res)
        # Display download button
        result = Image.fromarray(res)
        buffered = io.BytesIO()
        result.save(buffered, format="PNG")
        img_bytes = buffered.getvalue()
        st.download_button(label='Download Output', data=img_bytes, file_name='bilateral_blur_output.png', mime='image/png')

    elif option == 'Image Inpaint':

        stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 5)
        h, w = image.shape[:2]
        if w > 800:
            h_, w_ = int(h * 800 / w), 800
        else:
            h_, w_ = h, w

        # Create a canvas component.
        st.subheader("Draw over the areas you want to inpaint:")
        canvas_result = st_canvas(
            fill_color='white',
            stroke_width=stroke_width,
            stroke_color='black',
            background_image=Image.open(uploaded_file).resize((h_, w_)),
            update_streamlit=True,
            height=h_,
            width=w_,
            drawing_mode='freedraw',
            key="canvas",
        )
        stroke = canvas_result.image_data

        if stroke is not None:

            if st.sidebar.checkbox('show mask'):
                st.subheader("Mask")
                st.image(stroke)

            mask = cv2.split(stroke)[3]
            mask = np.uint8(mask)
            mask = cv2.resize(mask, (w, h))

        st.sidebar.caption('Happy with the selection?')
        option = st.sidebar.selectbox('Mode', ['None', 'Telea', 'NS', 'Compare both'])

        if option == 'Telea':
            st.subheader('Result of Telea')
            res = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)[:,:,::-1]
            st.image(res)
            # Display download button
            result = Image.fromarray(res)
            buffered = io.BytesIO()
            result.save(buffered, format="PNG")
            img_bytes = buffered.getvalue()
            st.download_button(label='Download Output', data=img_bytes, file_name='inpaint_telea_output.png', mime='image/png')
        elif option == 'Compare both':
            col1, col2 = st.columns(2)
            res1 = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)[:,:,::-1]
            res2 = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_NS)[:,:,::-1]
            with col1:
                st.subheader('Result of Telea')
                st.image(res1)
                # Display download button
                result1 = Image.fromarray(res1)
                buffered1 = io.BytesIO()
                result1.save(buffered1, format="PNG")
                img_bytes1 = buffered1.getvalue()
                st.download_button(label='Download Output', data=img_bytes1, file_name='inpaint_telea_output.png', mime='image/png')
            with col2:
                st.subheader('Result of NS')
                st.image(res2)
                # Display download button
                result2 = Image.fromarray(res2)
                buffered2 = io.BytesIO()
                result2.save(buffered2, format="PNG")
                img_bytes2 = buffered2.getvalue()
                st.download_button(label='Download Output', data=img_bytes2, file_name='inpaint_ns_output.png', mime='image/png')
        elif option == 'NS':
            st.subheader('Result of NS')
            res = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_NS)[:,:,::-1]
            st.image(res)
            # Display download button
            result = Image.fromarray(res)
            buffered = io.BytesIO()
            result.save(buffered, format="PNG")
            img_bytes = buffered.getvalue()
            st.download_button(label='Download Output', data=img_bytes, file_name='inpaint_ns_output.png', mime='image/png')
        else:
            pass