"""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')


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

if uploaded_file is not None:
    # Debug: Print uploaded file information
    # st.write("Uploaded file:", uploaded_file.name)

    # Convert the file to an opencv image.
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, 1)
    # Debug: Print image shape
    # st.write("Image shape:", image.shape)

    # Display the uploaded image immediately
    # st.image(image[:,:,::-1], caption='Uploaded Image')
    # 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.image(res)
    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.image(res)

    elif option == 'Image Inpaint':
        # Debug: Print selected option
        # st.write("Selected option for inpainting:", option)

        stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 5)
        # st.write("Stroke width:", stroke_width)  # Debug: Print stroke width

        h, w = image.shape[:2]
        # st.write("Original image dimensions (h, w):", h, w)  # Debug: Print dimensions
        if w > 800:
            h_, w_ = int(h * 800 / w), 800
        else:
            h_, w_ = h, w
        # st.write("Updated image dimensions (h_, w_):", h_, w_)  # Debug: Print dimensions

        # Create a canvas component.
        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",
        )
        
        # Debug: Print canvas result
        # st.write("Canvas result:", canvas_result)

        stroke = canvas_result.image_data
        if stroke is not None:
            # Debug: Print stroke data
            # st.write("Stroke data shape:", stroke.shape)

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

            mask = cv2.split(stroke)[3]
            mask = np.uint8(mask)
            mask = cv2.resize(mask, (w, h))
            # Debug: Print mask shape
            # st.write("Mask shape:", mask.shape)

        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)
            # Debug: Print result shape
            # st.write("Telea result shape:", res.shape)
        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)
            with col2:
                st.subheader('Result of NS')
                st.image(res2)
            if res1 is not None:
                # Display link.
                result1 = Image.fromarray(res1)
                st.sidebar.markdown(
                    get_image_download_link(result1, 'telea.png', 'Download Output of Telea'),
                    unsafe_allow_html=True)
            if res2 is not None:
                # Display link.
                result2 = Image.fromarray(res2)
                st.sidebar.markdown(
                    get_image_download_link(result2, 'ns.png', 'Download Output of NS'),
                    unsafe_allow_html=True)

        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)
        else:
            pass

    if res is not None:
        # Debug: Print final result shape
        # st.write("Final result shape:", res.shape)
        # Display link.
        result = Image.fromarray(res)
        st.sidebar.markdown(
            get_image_download_link(result, 'output.png', 'Download Output'),
            unsafe_allow_html=True)