import gradio as gr
import cv2
import os
import tempfile
import numpy as np
from utils import *
from algorithm import *

def make_video(video_path, outdir='./summarized_video',encoder='Kmeans'):
    if encoder not in ["Kmeans", "Sum of Squared Difference 01", "Sum of Squared Difference 02"]:
        encoder = "Kmeans"
    # nen them vao cac truong hop mo hinh khac
    margin_width = 50

    model, processor, device = load_model()

    # total_params = sum(param.numel() for param in model.parameters())
    # print('Total parameters: {:.2f}M'.format(total_params / 1e6))

    if os.path.isfile(video_path):
        if video_path.endswith('txt'):
            with open(video_path, 'r') as f:
                lines = f.read().splitlines()
        else:
            filenames = [video_path]
    else:
        filenames = os.listdir(video_path)
        filenames = [os.path.join(video_path, filename) for filename in filenames if not filename.startswith('.')]
        filenames.sort()
    
    for k, filename in enumerate(filenames):
        print('Progress {:}/{:},'.format(k+1, len(filenames)), 'Processing', filename)
        
        raw_video = cv2.VideoCapture(filename)
        frame_width, frame_height = int(raw_video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(raw_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frame_rate = int(raw_video.get(cv2.CAP_PROP_FPS))
        #length = int(raw_video.get(cv2.CAP_PROP_FRAME_COUNT))
        output_width = frame_width * 2 + margin_width
        
        filename = os.path.basename(filename)

        # Find the size to resize
        if "shortest_edge" in processor.size:
            height = width = processor.size["shortest_edge"]
        else:
            height = processor.size["height"]
            width = processor.size["width"] 
        resize_to = (height, width)

        # F/Fs
        clip_sample_rate = 1
        # F
        num_frames = 8

        frames = []
        features = []

        # output_path = os.path.join(outdir, filename[:filename.rfind('.')] + '_video_depth.mp4')
        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as tmpfile:
            output_path = tmpfile.name
        #out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"avc1"), frame_rate, (output_width, frame_height))
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, frame_rate, (output_width, frame_height))
        # count=0
        
        while raw_video.isOpened():
            ret, raw_frame = raw_video.read()
            if not ret:
                break
            
            raw_frame = cv2.resize(raw_frame, resize_to)
            frames.append(raw_frame)

        # Find key frames by selecting frames with clip_sample_rate
        key_frames = frames[::clip_sample_rate] 
        #print('total of frames after sample:', len(selected_frames))

        # Remove redundant frames to make the number of frames can be divided by num_frames
        num_redudant_frames = len(key_frames) - (len(key_frames) % num_frames)

        # Final key frames
        final_key_frames = key_frames[:num_redudant_frames]
        #print('total of frames after remove redundant frames:', len(selected_frames))

        for i in range(0, len(final_key_frames), num_frames):
            if i % num_frames*50 == 0:
                print(f"Loading {i}/{len(final_key_frames)}")
        
        # Input clip to the model
        input_frames = final_key_frames[i:i+num_frames]
        # Extract features
        batch_features = extract_features(input_frames, device, model, processor)
        # Convert to numpy array to decrease the memory usage
        batch_features = np.array(batch_features.cpu().detach().numpy())
        features.extend(batch_features)

        number_of_clusters = round(len(features)*0.15)

        selected_frames = []
        if encoder == "Kmeans":
            selected_frames = kmeans(features, number_of_clusters)
        elif encoder == "Sum of Squared Difference 01":
            selected_frames = tt01(features, 400)
        else:
            selected_frames = tt02(features, 400)

        video_writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), frame_rate, (frames[0].shape[1], frames[0].shape[0]))
        for idx in selected_frames:
            video_writer.write(frames[idx])
        
        raw_video.release()
        video_writer.release()
        print("Completed summarizing the video (wait for a moment to load).")
        return output_path

css = """
#img-display-container {
    max-height: 100vh;
    }
#img-display-input {
    max-height: 80vh;
    }
#img-display-output {
    max-height: 80vh;
    }
"""

title = "# Video Summarization Demo"
description = """Video Summarization using Timesformer.

Author: Nguyen Hoai Nam.
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    gr.Markdown("### Video Summarization demo")

    with gr.Row():
        input_video = gr.Video(label="Input Video")
        algorithm_type = gr.Dropdown(["Kmeans", "Sum of Squared Difference 01", "Sum of Squared Difference 02"], type="value", label='Algorithm')
    submit = gr.Button("Submit")
    processed_video = gr.Video(label="Summarized Video")

    def on_submit(uploaded_video,algorithm_type):
                
        # Process the video and get the path of the output video
        #output_video_path = make_video(uploaded_video,encoder=model_type)
        pass
        #return output_video_path

    submit.click(on_submit, inputs=[input_video, algorithm_type], outputs=processed_video)

    #example_files = os.listdir('assets/examples_video')
    #example_files.sort()
    #example_files = [os.path.join('assets/examples_video', filename) for filename in example_files]
    #examples = gr.Examples(examples=example_files, inputs=[input_video], outputs=processed_video, fn=on_submit, cache_examples=True)
    

if __name__ == '__main__':
    demo.queue().launch()