# App code based on:
# Model based on:

import numpy as np
import pandas as pd
import streamlit as st
import os
from datetime import datetime
from PIL import Image
from streamlit_drawable_canvas import st_canvas
from io import BytesIO
from copy import deepcopy

from src.core import process_inpaint


st.title("AI Photo Colorization")

st.image(open("assets/demo.png", "rb").read())

st.markdown(
    """
    Colorizing black & white photo can be expensive and time consuming. We introduce AI that can colorize
    grayscale photo in seconds. **Just upload your grayscale image, then click colorize.**
    """
)
uploaded_file = st.file_uploader("Choose image", accept_multiple_files=False, type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    bytes_data = uploaded_file.getvalue()
    img_input = Image.open(BytesIO(bytes_data)).convert("RGBA")

    if uploaded_file is not None and st.button("Colorize!"):
        
        with st.spinner("AI is doing the magic!"):
            img_output = """TODO"""
        
        # NOTE: Calm! I'm not logging the input and outputs.
        # It is impossible to access the filesystem in spaces environment.
        now = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
        img_input.convert("RGB").save(f"./output/{now}.jpg")
        Image.fromarray(img_output).convert("RGB").save(f"./output/{now}-edited.jpg")
        
        st.write("AI has finished the job!")
        st.image(img_output)
        # reuse = st.button('Edit again (Re-use this image)', on_click=set_image, args=(inpainted_img, ))
        
        with open(f"./output/{now}-edited.jpg", "rb") as fs:
            uploaded_name = os.path.splitext(uploaded_file.name)[0]
            st.download_button(
                label="Download",
                data=fs,
                file_name=f'edited_{uploaded_name}.jpg',
            )
            
        st.info("**TIP**: If the result is not perfect, you can download then "
                "re-upload the result then remove the artifacts.")