import math
import time
import numpy as np
import streamlit as st
from PIL import Image
import cv2

hide_streamlit_style = """
<style>
    #root > div:nth-child(1) > div > div > div > div > section > div {padding-top: 0rem;
    padding-left: 1%;
    }
</style>

"""

st.set_page_config(layout="wide")



@st.cache_resource(show_spinner=False)
def loadModel(n):
    super_res = cv2.dnn_superres.DnnSuperResImpl_create()
    super_res.readModel('models/ESPCN_x'+n+'.pb')
    return super_res

# on removing (show_spinner=False), it will show that fuction is running on web app
@st.cache_data(show_spinner=False)
def upscale(file,task):
    with open(file.name, "wb") as f:
        f.write(file.getbuffer())
        print('No file found, so added in list files')
    if isinstance(task,str):
        super_res = loadModel(task)
        super_res.setModel('espcn', int(task))
        if file.type.split('/')[0] == 'image':
            img = cv2.imread(file.name)
            upscaled_image = super_res.upsample(img)
            print('I upscaled upto',task,'times')
            cv2.imwrite("processed_"+file.name,upscaled_image)
        return True
    else:
        req_width,req_height = int(task[0]),int(task[1])
        if file.type.split('/')[0] == 'image':
            img = cv2.imread(file.name)
            actual_width,actual_height = img.shape[1],img.shape[0]
            w_ratio,h_ratio = req_width/actual_width , req_height/actual_height
            if min([w_ratio,h_ratio]) <= 1.0:
                img = cv2.resize(img,(req_width,req_height))
                print("I did resizing only!")
                cv2.imwrite("processed_" + file.name, img)
                return True
            # rounding off the ratios
            w_ratio,h_ratio = math.ceil(w_ratio),math.ceil(h_ratio)
            # find bigger number
            upscale_number = max(w_ratio,h_ratio)

            # task can be greater than 4 but we can upscale upto 4. So setting task to 4.
            if upscale_number >= 4:
                upscale_number = 4

            super_res = loadModel(str(upscale_number))
            super_res.setModel('espcn', int(upscale_number))
            upscaled_image = super_res.upsample(img)
            print("Before resizing ",(upscaled_image.shape[1], upscaled_image.shape[0]))
            upscaled_image = cv2.resize(upscaled_image,(task[0],task[1]))
            print("Final size got: ",(upscaled_image.shape[1],upscaled_image.shape[0]))

            print("I upscale upto", upscale_number , "times and then resize it.")

            cv2.imwrite("processed_" + file.name, upscaled_image)

            return True

        return "It's second"


if 'disable_opt2' not in st.session_state:
    st.session_state.disable_opt2 = True
if 'disable_opt1' not in st.session_state:
    st.session_state.disable_opt1 = False
if 'disable_download' not in st.session_state:
    st.session_state.disable_download = True


st.markdown(hide_streamlit_style, unsafe_allow_html=True)

col1,_,col2 = st.columns([6,1,3],gap="small")

def toggle_state_opt1():

    if st.session_state.get("opt1") == True:
        st.session_state.opt2 = False
        st.session_state.disable_opt2 = True

    else:
        st.session_state.opt2 = True
        st.session_state.disable_opt2 = False

def toggle_state_opt2():
    if st.session_state.get("opt2") == True:
        st.session_state.opt1 = False
        st.session_state.disable_opt1 = True
    else:
        st.session_state.opt1 = True
        st.session_state.disable_opt1 = False

# Update the states based on user selection before drawing the widgets in the web page
toggle_state_opt2()
toggle_state_opt1()

with col1:
    file = st.file_uploader("",type=['png','jpeg','jpg','pgm','jpe'])
    if file is not None:
        # writing file and saving its details in dict for further processing



        if file.type.split('/')[0] == "image":
            image = Image.open(file)
            st.image(image,caption="Upload Image", use_column_width=True)
        elif file.type.split('/')[0] == 'video':
            st.video(file)


with col2:
    st.markdown("\n")
    st.markdown("\n")
    st.markdown("\n")

    st.subheader(" UPSCALE RESOLUTION UP TO")
    st.markdown("\n")
    st.markdown("\n")

    opt1 = st.checkbox("MULTIPLES OF",key="opt1",value=True,on_change=toggle_state_opt1)
    st.selectbox("SELECT", ["2", "3", "4"],key="opt1_selBox",disabled=st.session_state.disable_opt1)

    st.markdown("\n")
    st.markdown("\n")
    opt2 = st.checkbox("CUSTOM SIZE",key="opt2",on_change=toggle_state_opt2)

    subCol3, subCol4 = st.columns(2)

    with subCol3:
        st.number_input("Width", step=1, min_value=150,max_value=3840, value=900, key="width",disabled=st.session_state.disable_opt2)

    with subCol4:
        st.number_input("Height", step=1, min_value=150,max_value=2160, value=900, key="height",disabled=st.session_state.disable_opt2)

    st.markdown("\n")
    st.markdown("\n")

    _, dcol, _ = st.columns([1,5,1],gap="small")

    with dcol:

        if st.button("PROCEED",use_container_width=True) and file is not None:
            if st.session_state.get('opt1') == True:
                task = st.session_state.opt1_selBox
            else:
                task = [st.session_state.width, st.session_state.height]
            print(task)
            st.session_state.disable_download = not upscale(file,task)


            #print(resulted_file.shape)

        st.markdown("\n")
        st.markdown("\n")

        if file is None:
            st.session_state.disable_download = True

        if st.session_state.disable_download == True:
            st.button("DOWNLOAD FILE",disabled=True,use_container_width=True)
        else:
            with open('processed_'+file.name, "rb") as download_file:
                st.download_button(label="Download image", data=download_file,
                                   file_name= 'processed_'+file.name, mime= "image/png",
                                   use_container_width=True, disabled=st.session_state.disable_download)



st.markdown("\n")
st.markdown("\n")
st.info("DESCRIPTION :   This web app is a free tool that allow user to upscale or resize media file resolution.")