import time
import streamlit as st
import subprocess
import numpy as np
from PIL import Image
from io import BytesIO
from models.HAT.hat import *
from models.RCAN.rcan import *
from models.SRGAN.srgan import *
from models.VDSR.vdsr import *
from models.Interpolation.nearest_neighbor import NearestNeighbor_for_deployment
from models.Interpolation.bilinear import Bilinear_for_deployment
from models.Interpolation.bicubic import Bicubic_for_deployment

subprocess.call('pip install natsort', shell=True)
from models.SRFlow.srflow import * 

# Initialize session state for enhanced images
if 'nearest_enhanced_image' not in st.session_state:
    st.session_state['nearest_enhanced_image'] = None
if 'bilinear_enhanced_image' not in st.session_state:
    st.session_state['bilinear_enhanced_image'] = None
if 'bicubic_enhanced_image' not in st.session_state:
    st.session_state['bicubic_enhanced_image'] = None
if 'hat_enhanced_image' not in st.session_state:
    st.session_state['hat_enhanced_image'] = None
if 'rcan_enhanced_image' not in st.session_state:
    st.session_state['rcan_enhanced_image'] = None
if 'srgan_enhanced_image' not in st.session_state:
    st.session_state['srgan_enhanced_image'] = None
if 'srflow_enhanced_image' not in st.session_state:
    st.session_state['srflow_enhanced_image'] = None
if 'vdsr_enhanced_image' not in st.session_state:
    st.session_state['vdsr_enhanced_image'] = None

# Initialize session state for button clicks
if 'nearest_clicked' not in st.session_state:
    st.session_state['nearest_clicked'] = False
if 'bilinear_clicked' not in st.session_state:
    st.session_state['bilinear_clicked'] = False
if 'bicubic_clicked' not in st.session_state:
    st.session_state['bicubic_clicked'] = False
if 'hat_clicked' not in st.session_state:
    st.session_state['hat_clicked'] = False
if 'rcan_clicked' not in st.session_state:
    st.session_state['rcan_clicked'] = False
if 'srgan_clicked' not in st.session_state:
    st.session_state['srgan_clicked'] = False 
if 'srflow_clicked' not in st.session_state:
    st.session_state['srflow_clicked'] = False
if 'vdsr_clicked' not in st.session_state:
    st.session_state['vdsr_clicked'] = False

st.markdown("<h1 style='text-align: center'>Image Super Resolution</h1>", unsafe_allow_html=True)

# Sidebar for navigation
st.sidebar.title("Options")
app_mode = st.sidebar.selectbox("Choose the input source", ["Upload image", "Take a photo"])

# Depending on the choice, show the uploader widget or webcam capture
if app_mode == "Upload image":
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png"], on_change=lambda: reset_states())
    if uploaded_file is not None:
        image = Image.open(uploaded_file).convert("RGB")
elif app_mode == "Take a photo":
    camera_input = st.camera_input("Take a picture", on_change=lambda: reset_states())
    if camera_input is not None:
        image = Image.open(camera_input).convert("RGB")
        
def reset_states():
    st.session_state['hat_enhanced_image'] = None
    st.session_state['rcan_enhanced_image'] = None
    st.session_state['srgan_enhanced_image'] = None
    st.session_state['srflow_enhanced_image'] = None
    st.session_state['bicubic_enhanced_image'] = None
    st.session_state['bilinear_enhanced_image'] = None
    st.session_state['nearest_enhanced_image'] = None
    st.session_state['vdsr_enhanced_image'] = None
    st.session_state['hat_clicked'] = False
    st.session_state['rcan_clicked'] = False
    st.session_state['srgan_clicked'] = False
    st.session_state['srflow_clicked'] = False
    st.session_state['bicubic_clicked'] = False
    st.session_state['bilinear_clicked'] = False
    st.session_state['nearest_clicked'] = False
    st.session_state['vdsr_clicked'] = False
    
def get_image_download_link(img, filename):
    """Generates a link allowing the PIL image to be downloaded"""
    # Convert the PIL image to Bytes
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    return st.download_button(
        label="Download Image",
        data=buffered.getvalue(),
        file_name=filename,
        mime="image/png"
    )
    
if 'image' in locals():
    # st.image(image, caption='Uploaded Image', use_column_width=True)
    st.write("")
    # ------------------------ Nearest Neighbor ------------------------ #
    if st.button('Enhance with Nearest Neighbor'):
        with st.spinner('Processing using Nearest Neighbor...'):
            enhanced_image = NearestNeighbor_for_deployment(image)
            st.session_state['nearest_enhanced_image'] = enhanced_image
            st.session_state['nearest_clicked'] = True
            st.success('Done!')
    if st.session_state['nearest_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['nearest_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['nearest_enhanced_image'], 'nearest_enhanced.jpg')

    # ------------------------ Bilinear ------------------------ #
    if st.button('Enhance with Bilinear'):
        with st.spinner('Processing using Bilinear...'):
            enhanced_image = Bilinear_for_deployment(image)
            st.session_state['bilinear_enhanced_image'] = enhanced_image
            st.session_state['bilinear_clicked'] = True
            st.success('Done!')
    if st.session_state['bilinear_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['bilinear_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['bilinear_enhanced_image'], 'bilinear_enhanced.jpg')
    
    # ------------------------ Bicubic ------------------------ #
    if st.button('Enhance with Bicubic'):
        with st.spinner('Processing using Bicubic...'):
            enhanced_image = Bicubic_for_deployment(image)
            st.session_state['bicubic_enhanced_image'] = enhanced_image
            st.session_state['bicubic_clicked'] = True
            st.success('Done!')
    if st.session_state['bicubic_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['bicubic_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['bicubic_enhanced_image'], 'bicubic_enhanced.jpg')

    # ------------------------ HAT ------------------------ #
    if st.button('Enhance with HAT'):
        with st.spinner('Processing using HAT...'):                             
            with st.spinner('Wait for it... the model is processing the image'):                                                  
                enhanced_image = HAT_for_deployment(image)
                st.session_state['hat_enhanced_image'] = enhanced_image
            st.session_state['hat_clicked'] = True
            st.success('Done!')
    if st.session_state['hat_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['hat_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['hat_enhanced_image'], 'hat_enhanced.jpg')
    
    # ------------------------ RCAN ------------------------ #
    if st.button('Enhance with RCAN'):
        with st.spinner('Processing using RCAN...'):
            with st.spinner('Wait for it... the model is processing the image'):
                rcan_model = RCAN()
                device = torch.device('cpu') if not torch.cuda.is_available() else torch.device('cuda')
                rcan_model.load_state_dict(torch.load('models/RCAN/rcan_checkpoint.pth', map_location=device))
                enhanced_image = rcan_model.inference(image)
                st.session_state['rcan_enhanced_image'] = enhanced_image
            st.session_state['rcan_clicked'] = True 
            st.success('Done!')
    if st.session_state['rcan_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['rcan_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['rcan_enhanced_image'], 'rcan_enhanced.jpg')
    
    # --------------------------SRGAN-------------------------- #
    if st.button('Enhance with SRGAN'):
        with st.spinner('Processing using SRGAN...'):
            with st.spinner('Wait for it... the model is processing the image'):
                srgan_model = GeneratorResnet()
                device = torch.device('cpu') if not torch.cuda.is_available() else torch.device('cuda')
                srgan_model = torch.load('models/SRGAN/srgan_checkpoint.pth', map_location=device)
                enhanced_image = srgan_model.inference(image)
                st.session_state['srgan_enhanced_image'] = enhanced_image
            st.session_state['srgan_clicked'] = True 
            st.success('Done!')
    if st.session_state['srgan_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['srgan_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['srgan_enhanced_image'], 'srgan_enhanced.jpg')
    
    # ------------------------ SRFlow ------------------------ #
    if st.button('Enhance with SRFlow'):
        with st.spinner('Processing using SRFlow...'):
            with st.spinner('Wait for it... the model is processing the image'):
                enhanced_image = return_SRFlow_result(image)
                st.session_state['srflow_enhanced_image'] = enhanced_image        
            st.session_state['srflow_clicked'] = True 
            st.success('Done!')
    if st.session_state['srflow_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['srflow_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['srflow_enhanced_image'], 'srflow_enhanced.jpg')

    # ------------------------ VDSR ------------------------ #
    if st.button('Enhance with VDSR'):
        with st.spinner('Processing using VDSR...'):
            # Load the VDSR model
            vdsr_model = torch.load('models/VDSR/vdsr_checkpoint.pth', map_location=torch.device('cpu'))
            enhanced_image = vdsr_model.inference(image)
            st.session_state['vdsr_enhanced_image'] = enhanced_image
            st.session_state['vdsr_clicked'] = True
            st.success('Done!')
    if st.session_state['vdsr_enhanced_image'] is not None:
        col1, col2 = st.columns(2)
        col1.header("Original")
        col1.image(image, use_column_width=True)
        col2.header("Enhanced")
        col2.image(st.session_state['vdsr_enhanced_image'], use_column_width=True)
        with col2:
            get_image_download_link(st.session_state['vdsr_enhanced_image'], 'vdsr_enhanced.jpg')