'''
python colorize_video.py xiangling_skt.mp4 xiangling_clr.mp4 --model ECCV16
'''

import os
import tempfile
import time
import argparse

import cv2
import moviepy.editor as mp
import numpy as np
from tqdm import tqdm

from models.deep_colorization.colorizers import eccv16
from utils import colorize_frame, change_model

def colorize_video(input_path, output_path, model_name):
    # Load the model
    loaded_model = change_model("None", model_name)

    # Open the video using cv2.VideoCapture
    video = cv2.VideoCapture(input_path)

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

    # Colorize video frames and store in a list
    output_frames = []

    for _ 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))

    # Merge frames to video
    frame_size = output_frames[0].shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for MP4 video
    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_size[1], frame_size[0]))

    for frame in output_frames:
        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        out.write(frame_bgr)

    out.release()

    # Close the video capture
    video.release()
    
    # Convert the output video to a format compatible with Streamlit
    clip = mp.VideoFileClip(output_path)

    # Write the video file
    #clip.write_videofile(output_path, codec="libx264", audio=audio is not None)
    clip.write_videofile(output_path)
    
    print(f"Colorized video saved to {output_path}")

def main():
    parser = argparse.ArgumentParser(description="Colorize a black and white video.")
    parser.add_argument("input_path", type=str, help="Path to the input video file.")
    parser.add_argument("output_path", type=str, help="Path to save the colorized video file.")
    parser.add_argument("--model", type=str, default="ECCV16", choices=["ECCV16", "SIGGRAPH17"], help="Model to use for colorization.")

    args = parser.parse_args()

    colorize_video(args.input_path, args.output_path, args.model)

if __name__ == "__main__":
    main()