import string

import random

import numpy as np

import face_recognition

from scipy.optimize import linear_sum_assignment

import pyspark.sql.types as T

import pyspark.sql.functions as F

from pyspark.sql import DataFrame, Row

from pyspark.sql.window import Window

from rikai.types.geometry import Box2d

from rikai.types import YouTubeVideo, Segment

from rikai.spark.types import YouTubeVideoType

from rikai.spark.types.geometry import Box2dType

from rikai.spark.functions import video_to_images

# Sample dataframe of youtube clips

df = spark.createDataFrame(

        [

            ("nf8ySuesAPg", YouTubeVideo(vid="nf8ySuesAPg"), Segment(30, 120))

        ],

        ["vid", "youtube_video", "segment"],

    )

df = df.withColumn("video_frames", video_to_images(F.col("youtube_video"), F.col("segment"), F.lit(5))) \

          .withColumn("frame", F.explode(F.col("video_frames"))) \

          .drop("video_frames")

display(df)

def bbox_helper(bbox):

  top, right, bottom, left = bbox

  bbox = [top, left, bottom, right]

  return list(map(lambda x: max(x, 0), bbox))

@udf(returnType=T.ArrayType(Box2dType()))

def face_detector(image):

  img = image.to_numpy()

  faces = face_recognition.face_locations(img)

  return [Box2d(*bbox_helper(f)) for f in faces]

df = df.withColumn("faces", face_detector(F.col("frame")))

display(df)

df = df.select("vid", "youtube_video", "frame", "faces") \

       .withColumn("detections", F.explode("faces")) \

       .withColumn("detect_id", F.lit("")) \

       .withColumn("annot", F.struct([F.col("detections"), F.col("detect_id")])) \

       .groupBy("vid", "youtube_video", "frame").agg(F.collect_list("annot").alias("annots"))

display(df)