BATCH_SIZE = 64
DOWNSAMPLE = 24
FOLDER_PATH = "."

import phash_jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from PIL import Image
import statistics
from decord import VideoReader
from decord import cpu
import gradio

def binary_array_to_hex(arr):
	"""
	Function to make a hex string out of a binary array.
	"""
	bit_string = ''.join(str(b) for b in 1 * arr.flatten())
	width = int(jnp.ceil(len(bit_string) / 4))
	return '{:0>{width}x}'.format(int(bit_string, 2), width=width)

def compute_batch_hashes(vid_path):
  kwargs={"width": 64, "height":64}
  vr = VideoReader(vid_path, ctx=cpu(0), **kwargs)
  hashes = []
  h_prev = None
  batch = []
  for i in range(0, len(vr), DOWNSAMPLE * BATCH_SIZE):
      print(f"batch_{i}")
      ids = [id for id in range(i, min(i + DOWNSAMPLE * BATCH_SIZE, len(vr)), DOWNSAMPLE)]
      vr.seek(0)
      batch = jnp.array(vr.get_batch(ids).asnumpy())
      batch_h =  phash_jax.batch_phash(batch)
      for i in range(len(ids)):
        h = batch_h[i]
        if h_prev == None:
          h_prev=h
        hashes.append({"frame_id":ids[i], "hash": binary_array_to_hex(h), "distance": int(phash_jax.hash_dist(h, h_prev))})
        h_prev = h
  return gradio.update(value=hashes, visible=False)

def plot_hash_distance(hashes, threshold):
  fig = plt.figure()
  ids = [h["frame_id"] for h in hashes]
  distances = [h["distance"] for h in hashes]
  plt.plot(ids, distances, ".")
  plt.plot(ids, [threshold]* len(ids), "r-")
  return fig

def compute_threshold(hashes):
  min_length = 24 * 3
  ids = [h["frame_id"] for h in hashes]
  distances = [h["distance"] for h in hashes]
  thrs_ = sorted(list(set(distances)),reverse=True)
  best = thrs_[0] - 1
  for threshold in thrs_[1:]:
    durations = []
    i_start=0
    for i, h in enumerate(hashes):
      if h["distance"] > threshold and hashes[i-1]["frame_id"] -  hashes[i_start]["frame_id"] > min_length:
        durations.append(hashes[i-1]["frame_id"] - hashes[i_start]["frame_id"])
        i_start=i
    if len(durations) < (len(hashes) * DOWNSAMPLE / 24) / 20:
      best = threshold
  return best

def get_slides(vid_path, hashes, threshold):
    min_length = 24 * 1.5
    vr = VideoReader(vid_path, ctx=cpu(0))
    slideshow = []
    i_start = 0
    for i, h in enumerate(hashes):
        if h["distance"] > threshold and hashes[i-1]["frame_id"] -  hashes[i_start]["frame_id"] > min_length:
            path=f'{FOLDER_PATH}/{vid_path.split("/")[-1].split(".")[0]}_{i_start}_{i-1}.png'
            Image.fromarray(vr[hashes[i-1]["frame_id"]].asnumpy()).save(path)
            slideshow.append({"slide": path, "start": i_start, "end": i-1})
            i_start=i
    path=f'{FOLDER_PATH}/{vid_path.split("/")[-1].split(".")[0]}_{i_start}_{len(vr)-1}.png'
    Image.fromarray(vr[-1].asnumpy()).save(path)
    slideshow.append({"slide": path, "start": i_start, "end": len(vr)-1})
    return [s["slide"] for s in slideshow]

def trigger_plots(f2f_distance_plot, hashes, threshold):
  # if not hist_plot.get_config()["visible"] and len(hashes.get_config()["value"]) > 0 :
  return gradio.update(value=plot_hash_distance(hashes, threshold))

def set_visible():
  return gradio.update(visible=True)
  
demo = gradio.Blocks()

with demo:
    with gradio.Row():
        with gradio.Column():
          with gradio.Row():
            vid=gradio.Video(mirror_webcam=False)
          with gradio.Row():
            btn_vid_proc = gradio.Button("Compute hashes")
          with gradio.Row():
            hist_plot = gradio.Plot(label="Frame to frame hash distance histogram", visible=False)
        with gradio.Column():
            hashes = gradio.JSON()
            with gradio.Column(visible=False) as result_row:
              btn_plot = gradio.Button("Plot & compute optimal threshold")
              threshold = gradio.Slider(minimum=1, maximum=30, value=5, label="Threshold")
              f2f_distance_plot = gradio.Plot(label="Frame to frame hash distance")
              btn_slides = gradio.Button("Extract Slides")
    with gradio.Row():
      slideshow = gradio.Gallery(label="Extracted slides")
      slideshow.style(grid=6)
    btn_vid_proc.click(fn=compute_batch_hashes, inputs=[vid], outputs=[hashes])
    hashes.change(fn=set_visible, inputs=[], outputs=[result_row])
    btn_plot.click(fn=compute_threshold, inputs=[hashes], outputs=[threshold])
    btn_plot.click(fn=trigger_plots, inputs=[f2f_distance_plot, hashes, threshold], outputs=[f2f_distance_plot])
    threshold.change(fn=plot_hash_distance, inputs=[hashes, threshold], outputs=f2f_distance_plot)
    btn_slides.click(fn=get_slides, inputs=[vid, hashes, threshold], outputs=[slideshow])

demo.queue(default_enabled=True).launch()