import tempfile

import numpy as np
import gradio as gr
import tensorflow as tf
from moviepy.editor import VideoFileClip
from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
from ultralytics import YOLO

from core.model import load_classifier
from core.inference import FrameProcessor
print("Tensorflow version " + tf.__version__)

print('Load classifier.')
classifier_path = 'weights/classifier-8-epoch10.keras'
classifier = load_classifier(classifier_path)

print('Load detector.')
detector_path = 'weights/yolov8n.pt'
detector = YOLO(detector_path)

def fn(video: gr.Video):
  print('Process video.')
  with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as f:
    output = f.name
    clip = VideoFileClip(video)
    process_frame = FrameProcessor(detector, classifier)
    processed_frames = []
    for frame in clip.iter_frames():
      processed_frames.append(process_frame(frame))
      yield processed_frames[-1], None
    processed_clip = ImageSequenceClip(processed_frames, clip.fps)
    processed_clip.audio = clip.audio
    clip.write_videofile(output, fps=clip.fps, audio_codec='aac', logger=None)
  yield processed_frames[-1], output

inputs = gr.Video(sources=['upload'], label='Input Video')
outputs = [
  gr.Image(interactive=False, label='Last Frame Processed'),
  gr.Video(interactive=False, label='Aeroplane Position and Action Marked')]

examples = [
  ['examples/ZFLFDfovqls_001310_001320.mp4'], # cspell: disable-line
  ['examples/Zv7GyH-fpEY_2023.0_2033.0.mp4']]

iface = gr.Interface(
  title='Aeroplane Position and Action Detection',
  description='Detect aeroplane position and action in a video.',
  theme='soft',
  fn=fn,
  inputs=inputs,
  outputs=outputs,
  examples=examples,
  cache_examples=False)
iface.launch()