import streamlit as st
from streamlit_lottie import st_lottie
from models.deep_colorization.colorizers import *
from utils import load_lottieurl, format_time, colorize_frame
import os
import tempfile
import cv2
import moviepy.editor as mp
import time
from tqdm import tqdm

st.set_page_config(page_title="Image & Video Colorizer", page_icon="🎨", layout="wide")


loaded_model = eccv16(pretrained=True).eval()
current_model = "None"


def change_model(current_model, model):
    if current_model != model:
        if model == "ECCV16":
            loaded_model = eccv16(pretrained=True).eval()
        elif model == "SIGGRAPH17":
            loaded_model = siggraph17(pretrained=True).eval()
        return loaded_model
    else:
        raise Exception("Model is the same as the current one.")


col1, col2 = st.columns([1, 3])
with col1:
    lottie = load_lottieurl("https://assets5.lottiefiles.com/packages/lf20_RHdEuzVfEL.json")
    st_lottie(lottie)

with col2:
    st.write("""
    ## B&W Videos Colorizer
    ##### Upload a black and white video and get a colorized version of it.
    ###### I recommend starting with the first model and then experimenting with the second one.""")


def main():
    model = st.selectbox(
        "Select Model (Both models have their pros and cons, I recommend to try both and keep the best for your task)",
        ["ECCV16", "SIGGRAPH17"], index=0)

    loaded_model = change_model(current_model, model)
    st.write(f"Model is now {model}")

    uploaded_file = st.file_uploader("Upload your video here...", type=['mp4'])

    if st.button("Colorize"):
        if uploaded_file is not None:
            file_extension = os.path.splitext(uploaded_file.name)[1].lower()
            if file_extension == '.mp4':
                # Save the video file to a temporary location
                temp_file = tempfile.NamedTemporaryFile(delete=False)
                temp_file.write(uploaded_file.read())

                audio = mp.AudioFileClip(temp_file.name)

                # Open the video using cv2.VideoCapture
                video = cv2.VideoCapture(temp_file.name)

                # Get video information
                fps = video.get(cv2.CAP_PROP_FPS)

                col1, col2 = st.columns([0.5, 0.5])
                with col1:
                    st.markdown('<p style="text-align: center;">Before</p>', unsafe_allow_html=True)
                    st.video(temp_file.name)

                with col2:
                    st.markdown('<p style="text-align: center;">After</p>', unsafe_allow_html=True)

                    with st.spinner("Colorizing frames..."):
                        # Colorize video frames and store in a list
                        output_frames = []
                        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
                        progress_bar = st.empty()

                        start_time = time.time()
                        for i in tqdm(range(total_frames), unit='frame', desc="Progress"):
                            ret, frame = video.read()
                            if not ret:
                                break

                            colorized_frame = colorize_frame(frame, loaded_model)
                            output_frames.append((colorized_frame * 255).astype(np.uint8))

                            elapsed_time = time.time() - start_time
                            frames_completed = len(output_frames)
                            frames_remaining = total_frames - frames_completed
                            time_remaining = (frames_remaining / frames_completed) * elapsed_time

                            progress_bar.progress(frames_completed / total_frames)

                            if frames_completed < total_frames:
                                progress_bar.text(f"Time Remaining: {format_time(time_remaining)}")
                            else:
                                progress_bar.empty()

                    with st.spinner("Merging frames to video..."):
                        frame_size = output_frames[0].shape[:2]
                        output_filename = "output.mp4"
                        fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for MP4 video
                        out = cv2.VideoWriter(output_filename, fourcc, fps, (frame_size[1], frame_size[0]))

                        # Display the colorized video using st.video
                        for frame in output_frames:
                            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

                            out.write(frame_bgr)

                        out.release()

                        # Convert the output video to a format compatible with Streamlit
                        converted_filename = "converted_output.mp4"
                        clip = mp.VideoFileClip(output_filename)
                        clip = clip.set_audio(audio)

                        clip.write_videofile(converted_filename, codec="libx264")

                        # Display the converted video using st.video()
                        st.video(converted_filename)

                        # Add a download button for the colorized video
                        st.download_button(
                            label="Download Colorized Video",
                            data=open(converted_filename, "rb").read(),
                            file_name="colorized_video.mp4"
                        )

                        # Close and delete the temporary file after processing
                        video.release()
                        temp_file.close()


if __name__ == "__main__":
    main()
    st.markdown(
        "###### Made with :heart: by [Clément Delteil](https://www.linkedin.com/in/clementdelteil/) [![this is an image link](https://i.imgur.com/thJhzOO.png)](https://www.buymeacoffee.com/clementdelteil)")