# Dependencies, see also requirement.txt ;)
import gradio as gr
import cv2 
import numpy as np
import os

from scenedetect import open_video, SceneManager
from scenedetect.detectors import ContentDetector

from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip

# —————————————————————————————————————————————————

title = "Scene Edit Detection"
description = "Gradio demo of PyScene scenedetect, to automatically find every shots in a video sequence, give you timecode in/out for each shot, save each shot as a splitted mp4 video chunk for you to download, then display a thumbnail for each shot as a gallery output"

# —————————————————————————————————————————————————

# SET INPUTS
video_input = gr.Video(source="upload", format="mp4", label="Video Sequence", mirror_webcam=False)
threshold = gr.Slider(label="Threshold pixel comparison: if exceeded, triggers a scene cut. Default: 27.0", minimum=15.0, maximum=40.0, value=27.0)

# —————————————————————————————————————————————————

def convert_to_tuple(list):
    return tuple(list);


def find_scenes(video_path, threshold):
    # file name without extension
    filename = os.path.splitext(os.path.basename(video_path))[0]
    # Open our video, create a scene manager, and add a detector.
    video = open_video(video_path)
    scene_manager = SceneManager()
    scene_manager.add_detector(
        ContentDetector(threshold=threshold))
    
    # Start detection 
    scene_manager.detect_scenes(video, show_progress=True)
    scene_list = scene_manager.get_scene_list()
    
    # Push the list of scenes into data_outputs
    data_outputs.append(scene_list)
    gradio_components_outputs.append("json")
    #print(scene_list)
    
    timecodes = []
    timecodes.append({"title": filename + ".mp4", "fps": scene_list[0][0].get_framerate()})
    
    shots = []
    stills = []
    
    # For each shot found, set entry and exit points as seconds from frame number
    # Then split video into chunks and store them into shots List
    # Then extract first frame of each shot as thumbnail for the gallery
    for i, shot in enumerate(scene_list):
        
        # STEP 1
        # Get timecode in seconds
        framerate = shot[0].get_framerate()
        shot_in = shot[0].get_frames() / framerate
        shot_out = shot[1].get_frames() / framerate
        
        tc_in = shot[0].get_timecode()
        tc_out = shot[1].get_timecode()
        
        frame_in = shot[0].get_frames()
        frame_out = shot[1].get_frames()
        
        timecode = {"tc_in": tc_in, "tc_out": tc_out, "frame_in": frame_in, "frame_out": frame_out}
        timecodes.append(timecode)
        
        # Set name template for each shot
        target_name = "shot_" + str(i+1) + "_" + filename + ".mp4"
        
        # Split chunk
        ffmpeg_extract_subclip(video_path, shot_in, shot_out, targetname=target_name)
        
        # Push chunk into shots List
        shots.append(target_name)
        
        # Push each chunk into data_outputs
        data_outputs.append(target_name)
        gradio_components_outputs.append("video")
        
        # —————————————————————————————————————————————————
        
        # STEP 2
        # extract first frame of each shot with cv2
        vid = cv2.VideoCapture(video_path)
        fps = vid.get(cv2.CAP_PROP_FPS)
        print('frames per second =',fps)
        
        frame_id = shot[0].get_frames() # value from scene_list from step 1

        vid.set(cv2.CAP_PROP_POS_FRAMES, frame_id)
        ret, frame = vid.read()

        # Save frame as PNG file
        img = str(frame_id) + '_screenshot.png'
        cv2.imwrite(img,frame)
        
        # Push image into stills List
        stills.append(img)
    
    # Push the list of video shots into data_outputs for Gradio file component
    data_outputs.append(shots)
    gradio_components_outputs.append("file")
        
    # Push the list of still images into data_outputs
    data_outputs.append(stills)
    gradio_components_outputs.append("gallery")
    
    # This would have been used as gradio outputs, 
    # if we could set number of outputs after the interface launch
    # That's not (yet ?) possible
    results = convert_to_tuple(data_outputs)    
    print(results)
 
    # return List of shots as JSON, List of video chunks, List of still images  
    # *
    # Would be nice to be able to return my results tuple as outputs, 
    # while number of chunks found is not fixed:
    # return results 
    return timecodes, shots, stills

# —————————————————————————————————————————————————

# SET DATA AND COMPONENTS OUTPUTS

# This would be filled like this:
# data_outputs = [ [List from detection], "video_chunk_n0.mp4", "video_chunk_n1.mp4", ... , "video_chunk_n.mp4", [List of video filepath to download], [List of still images from each shot found] ]
data_outputs = []

# This would be filled like this:
# gradio_components_outputs = [ "json", "video", "video", ... , "video", "file", "gallery" ]
gradio_components_outputs = []


#SET OUTPUTS

# This would be nice if number of outputs could be set after Interface Launch: 
# because we do not know how many shots will be detected         
# gradio_components_outputs = [ "json", "video", "video", ... , "video", "file", "gallery" ]
# outputs = gradio_components_outputs

# ANOTHER SOLUTION WOULD BE USING A (FUTURE ?) "VIDEO GALLERY" GRADIO COMPONENT FROM LIST :)

outputs = [gr.JSON(label="Shots detected"), gr.File(label="Downloadable Shots"), gr.Gallery(label="Still Images from each shot").style(grid=3)]

# —————————————————————————————————————————————————
print('Hello Sylvain')
gr.Interface(fn=find_scenes, inputs=[video_input, threshold], outputs=outputs, title=title, description=description).launch()