#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright © 2017-2020 bily     Huazhong University of Science and Technology
#
# Distributed under terms of the MIT license.

"""Figure with horizontal scrollbar and play capabilities
For latest version, go to https://github.com/bilylee/videofig
Basic usage
-----------
Creates a figure with a horizontal scrollbar and shortcuts to scroll automatically.
The scroll range is 0 to NUM_FRAMES - 1. The function REDRAW_FUN(F, AXES) is called to
redraw at scroll position F (for example, REDRAW_FUNC can show the frame F of a video)
using AXES for drawing. F is an integer, AXES is a instance of [Axes class](https://matplotlib.org/api/axes_api.html)
This can be used not only to play and analyze standard videos, but it also lets you place
any custom Matplotlib plots and graphics on top.
The keyboard shortcuts are:
  Enter(Return) -- play/pause video (25 frames-per-second default).
  Backspace -- play/pause video 5 times slower.
  Right/left arrow keys -- advance/go back one frame.
  Page down/page up -- advance/go back 30 frames.
  Home/end -- go to first/last frame of video.
Advanced usage
--------------
videofig(NUM_FRAMES, REDRAW_FUNC, FPS, BIG_SCROLL)
Also specifies the speed of the play function (frames-per-second) and
the frame step of page up/page down (or empty for defaults).
videofig(NUM_FRAMES, REDRAW_FUNC, FPS, BIG_SCROLL, KEY_FUNC)
Also calls KEY_FUNC(KEY) with any keys that weren't processed, so you
can add more shortcut keys (or empty for none).
Example 1: Plot a dynamic sine wave
---------
  import numpy as np
  def redraw_fn(f, axes):
    amp = float(f) / 3000
    f0 = 3
    t = np.arange(0.0, 1.0, 0.001)
    s = amp * np.sin(2 * np.pi * f0 * t)
    if not redraw_fn.initialized:
      redraw_fn.l, = axes.plot(t, s, lw=2, color='red')
      redraw_fn.initialized = True
    else:
      redraw_fn.l.set_ydata(s)
  redraw_fn.initialized = False
  videofig(100, redraw_fn)
Example 2: Show images in a custom directory
---------
  import os
  import glob
  from scipy.misc import imread
  img_dir = 'YOUR-IMAGE-DIRECTORY'
  img_files = glob.glob(os.path.join(video_dir, '*.jpg'))
  def redraw_fn(f, axes):
    img_file = img_files[f]
    img = imread(img_file)
    if not redraw_fn.initialized:
      redraw_fn.im = axes.imshow(img, animated=True)
      redraw_fn.initialized = True
    else:
      redraw_fn.im.set_array(img)
  redraw_fn.initialized = False
  videofig(len(img_files), redraw_fn, play_fps=30)
Example 3: Show images together with object bounding boxes
----------
  import os
  import glob
  from scipy.misc import imread
  from matplotlib.pyplot import Rectangle
  video_dir = 'YOUR-VIDEO-DIRECTORY'
  img_files = glob.glob(os.path.join(video_dir, '*.jpg'))
  box_files = glob.glob(os.path.join(video_dir, '*.txt'))
  def redraw_fn(f, axes):
    img = imread(img_files[f])
    box = bbread(box_files[f])  # Define your own bounding box reading utility
    x, y, w, h = box
    if not redraw_fn.initialized:
      im = axes.imshow(img, animated=True)
      bb = Rectangle((x, y), w, h,
                     fill=False,  # remove background
                     edgecolor="red")
      axes.add_patch(bb)
      redraw_fn.im = im
      redraw_fn.bb = bb
      redraw_fn.initialized = True
    else:
      redraw_fn.im.set_array(img)
      redraw_fn.bb.set_xy((x, y))
      redraw_fn.bb.set_width(w)
      redraw_fn.bb.set_height(h)
  redraw_fn.initialized = False
  videofig(len(img_files), redraw_fn, play_fps=30)
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

import matplotlib
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.widgets import Slider
import matplotlib.gridspec as gridspec

__all__ = ['videofig']


def videofig(num_frames, redraw_func, play_fps=25, big_scroll=30, key_func=None,
             grid_specs=None, layout_specs=None, figname=None, save_dir=None, *args):
    """Figure with horizontal scrollbar and play capabilities
    This script is mainly inspired by the elegant work of João Filipe Henriques
      https://www.mathworks.com/matlabcentral/fileexchange/29544-figure-to-play-and-analyze-videos-with-custom-plots-on-top?focused=5172704&tab=function
    :param num_frames: an integer, number of frames in a sequence
    :param redraw_func: callable with signature redraw_func(f, axes)
                        used to draw a new frame at position f using axes, which is a instance of Axes class in matplotlib
    :param play_fps: an integer, number of frames per second, used to control the play speed
    :param big_scroll: an integer, big scroll number used when pressed page down or page up keys.
    :param key_func: optional callable which signature key_func(key), used to provide custom key shortcuts.
    :param grid_specs: optional dictionary, used to specify the gridspec of the main drawing pane.
                       For example, grid_specs = {'nrows': 2, 'ncols': 2} will create a gridspec with 2 rows and 2 cols.
    :param layout_specs: optional list, used to specify the layout of the gridspec of the main drawing pane.
                       For example, layout_specs = ['[:, 0]', '[:, 1]'] means:
                          gs = ... Some code to create the main drawing pane gridspec ...
                          ax1 = plt.subplot(gs[:, 0])
                          ax2 = plt.subplot(gs[:, 1])
    :param save_dir: a string, used to specify the directory to which figures will be saved.
    :param args: other optional arguments
    :return: None
    """
    # Check arguments
    check_int_scalar(num_frames, 'num_frames')
    check_callback(redraw_func, 'redraw_func')
    check_int_scalar(play_fps, 'play_fps')
    check_int_scalar(big_scroll, 'big_scroll')
    if key_func:
        check_callback(key_func, 'key_func')

    # Initialize figure
    if not figname:
        figname = 'VideoPlayer'
    fig_handle = plt.figure(num=figname)

    # We do not want to show the slider when saving figures
    slider_ratio = 3 if save_dir is None else 1e-3

    # We use GridSpec to support embedding multiple plots in the main drawing pane.
    # A nested grid demo can be found at https://matplotlib.org/users/plotting/examples/demo_gridspec06.py
    # Construct outer grid, which contains the main drawing pane and slider ball
    outer_grid = gridspec.GridSpec(2, 1,
                                   left=0, bottom=0, right=1, top=1,
                                   height_ratios=[97, slider_ratio], wspace=0.0, hspace=0.0)

    # Construct inner grid in main drawing pane
    if grid_specs is None:
        grid_specs = {'nrows': 1, 'ncols': 1}  # We will have only one axes by default

    inner_grid = gridspec.GridSpecFromSubplotSpec(subplot_spec=outer_grid[0], **grid_specs)
    if layout_specs:
        # eval() can't work properly in list comprehension which is inside a function.
        # Refer to: http://bugs.python.org/issue5242
        # Maybe I should find another way to implement layout_specs without using eval()...
        axes_handle = []
        for spec in layout_specs:
            axes_handle.append(plt.Subplot(fig_handle, eval('inner_grid' + spec)))
    else:
        num_inner_plots = grid_specs['nrows'] * grid_specs['ncols']
        axes_handle = [plt.Subplot(fig_handle, inner_grid[i]) for i in range(num_inner_plots)]

    for ax in axes_handle: fig_handle.add_subplot(ax)

    if len(axes_handle) == 1:
        axes_handle = axes_handle[0]
        axes_handle.set_axis_off()

    # Build scrollbar
    scroll_axes_handle = plt.Subplot(fig_handle, outer_grid[1])
    scroll_axes_handle.set_facecolor('lightgoldenrodyellow')
    fig_handle.add_subplot(scroll_axes_handle)
    scroll_handle = Slider(scroll_axes_handle, '', 0.0, num_frames - 1, valinit=0.0)

    def draw_new(_):
        redraw_func(int(scroll_handle.val), axes_handle)
        fig_handle.canvas.draw_idle()

    def scroll(new_f):
        new_f = min(max(new_f, 0), num_frames - 1)  # clip in the range of [0, num_frames - 1]
        cur_f = scroll_handle.val

        # Stop player at the end of the sequence
        if new_f == (num_frames - 1):
            play.running = False

        if cur_f != new_f:
            # move scroll bar to new position
            scroll_handle.set_val(new_f)

        return axes_handle

    def play(period):
        play.running ^= True  # Toggle state
        if play.running:
            frame_idxs = range(int(scroll_handle.val), num_frames)
            play.anim = FuncAnimation(fig_handle, scroll, frame_idxs,
                                      interval=1000 * period, repeat=False, blit=False)
            plt.draw()
        else:
            play.anim.event_source.stop()

    # Set initial player state
    play.running = False

    def key_press(event):
        key = event.key
        f = scroll_handle.val
        if key == 'left':
            scroll(f - 1)
        elif key == 'right':
            scroll(f + 1)
        elif key in ['pageup', 'up']:
            scroll(f - big_scroll)
        elif key in ['pagedown', 'down']:
            scroll(f + big_scroll)
        elif key == 'home':
            scroll(0)
        elif key == 'end':
            scroll(num_frames - 1)
        elif key == 'enter':
            play(1 / play_fps)
        elif key == 'backspace':
            play(5 / play_fps)
        else:
            if key_func:
                key_func(key)

    # Register events
    scroll_handle.on_changed(draw_new)
    fig_handle.canvas.mpl_connect('key_press_event', key_press)

    if save_dir is None:
        # Draw initial frame
        redraw_func(0, axes_handle)

        # Start playing
        play(1 / play_fps)

        # plt.show() has to be put in the end of the function,
        # otherwise, the program simply won't work, weird...
        plt.show()
    else:
        print('=> Saving plots to: {}...'.format(save_dir))
        os.makedirs(save_dir, exist_ok=True)

        # Due to the issue of https://github.com/matplotlib/matplotlib/issues/7614
        # we need to save the first image twice. Otherwise, the first saved image will be distorted.
        for index in [0] + list(range(0, num_frames)):
            redraw_func(index, axes_handle)
            fig_handle.canvas.draw_idle()
            if not isinstance(axes_handle, list):
                if matplotlib.get_backend() == 'nbAgg':
                    print('Warning: remove "%matplotlib notebook", otherwise saved images will be distorted!')
                extent = axes_handle.get_window_extent().transformed(fig_handle.dpi_scale_trans.inverted())
            else:
                extent = None
            fig_handle.savefig(os.path.join(save_dir, '{:04d}.jpg'.format(index)), bbox_inches=extent)
        print('=> Done')


def check_int_scalar(a, name):
    assert isinstance(a, int), '{} must be a int scalar, instead of {}'.format(name, type(name))


def check_callback(a, name):
    # Check http://stackoverflow.com/questions/624926/how-to-detect-whether-a-python-variable-is-a-function
    # for more details about python function type detection.
    assert callable(a), '{} must be callable, instead of {}'.format(name, type(name))


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--imgdir", type=str, default=None, help="directory containing JPEG images")
    args = parser.parse_args()

    if args.imgdir:
        import glob
        from scipy.misc import imread

        img_files = sorted(glob.glob(os.path.join(args.imgdir, '*.jpg')))


        def redraw_fn(f, axes):
            img_file = img_files[f]
            img = imread(img_file)
            if not redraw_fn.initialized:
                redraw_fn.im = axes.imshow(img, animated=True)
                redraw_fn.initialized = True
            else:
                redraw_fn.im.set_array(img)


        redraw_fn.initialized = False
        videofig(len(img_files), redraw_fn, play_fps=30)

    else:
        import numpy as np


        def redraw_fn(f, axes):
            amp = float(f) / 3000
            f0 = 3
            t = np.arange(0.0, 1.0, 0.001)
            s = amp * np.sin(2 * np.pi * f0 * t)
            if not redraw_fn.initialized:
                redraw_fn.l, = axes.plot(t, s, lw=2, color='red')
                redraw_fn.initialized = True
            else:
                redraw_fn.l.set_ydata(s)


        redraw_fn.initialized = False
        videofig(100, redraw_fn)
