#!/usr/bin/env python3
 
 """
 Zoom in on video motion.
 """
 
 import pyffstream
 import numpy as np
 import PIL.Image
 import cv2
 import scipy.signal
 
 
 def args_pre(parser):
     # Add arguments.
     parser.add_argument(
         '--margins', metavar=('L', 'R', 'T', 'B'), type=float, nargs=4,
         default=[0, 0, 0, 0],
         help="""
         margins (left, right, top, bottom, in percent) of output video
         (default: %(default)s)
         """)
     parser.add_argument(
         '--blur-factor', metavar='F', type=float, default=0.05,
         help="blur size factor (default: %(default)s)")
     parser.add_argument(
         '--blur-threshold', metavar='T', type=int, default=32,
         help="blur threshold (default: %(default)s)")
     parser.add_argument(
         '--lowpass-factor', metavar='F', type=float, default=0.00015,
         help="low-pass filter cutoff frequency factor (default: %(default)s)")
 
 
 def init(args):
     # Set arguments.
     args.history = int(30 * args.output_fps)
     args.blur_size = (
         int(np.ceil(args.working_width * args.blur_factor)) // 2 * 2 + 1
     )
     args.lost_size = args.working_width * args.working_height * 0.001
     args.b, args.a = scipy.signal.butter(
         1, args.working_width / args.output_fps * args.lowpass_factor
     )
     args.history_track_ratio = 0.025
     args.resample_ratios = [
         # (1.0, PIL.Image.NEAREST),
         (0.5, PIL.Image.BILINEAR),
         (0.2, PIL.Image.BICUBIC),
         (0.0, PIL.Image.LANCZOS),
     ]
 
     # Set state.
     class State:
         pass
     state = State()
     state.background_subtractor = cv2.createBackgroundSubtractorMOG2(
         history=args.history
     )
     state.filter_state = []
     return state
 
 
 def process(args, state, frame, frame_num):
     # Create debug frame.
     if args.debug:
         debug_frame = frame.copy()
     else:
         debug_frame = None
 
     # Subtract background, blur and threshold.
     foreground = state.background_subtractor.apply(frame)
     mask = cv2.compare(
         cv2.GaussianBlur(foreground, (args.blur_size, args.blur_size), 0),
         args.blur_threshold,
         cv2.CMP_GE,
     )
     if args.debug:
         debug_frame[mask > 0] = (0, 255, 0)
         debug_frame[foreground > 0] = (255, 0, 0)
 
     # Nothing interesting?
     if np.count_nonzero(mask) < args.lost_size:
         # Reset rectangle.
         x, y, w, h = 0, 0, args.working_width, args.working_height
     else:
         # Find bounding rectangle.
         x, y, w, h = cv2.boundingRect(mask)
     if args.debug:
         cv2.rectangle(
             debug_frame, (x, y), (x+w, y+h), (0, 255, 0), 2 * args.thickness
         )
 
     # Add rectangle margins.
     ml, mr, mt, mb = args.margins
     m = max(w, h)
     x = max(x - int(m * ml / 100), 0)
     y = max(y - int(m * mt / 100), 0)
     w = min(w + int(m * (ml+mr) / 100), args.working_width - x)
     h = min(h + int(m * (mt+mb) / 100), args.working_height - y)
     if args.debug:
         cv2.rectangle(
             debug_frame, (x, y), (x+w, y+h), (0, 0, 255), 2 * args.thickness
         )
 
     # Filter rectangle.
     x1, y1, x2, y2 = x, y, x+w, y+h
     if frame_num == args.start_frame:
         state.filter_state = [
             coord * scipy.signal.lfilter_zi(args.b, args.a)
             for coord in (x1, y1, x2, y2)
         ]
     (x1, y1, x2, y2), filter_state_next = zip(*(
         scipy.signal.lfilter(args.b, args.a, [coord], zi=zi)
         for coord, zi in
         zip((x1, y1, x2, y2), state.filter_state)
     ))
     if frame_num >= args.start_frame + args.history * args.history_track_ratio:
         state.filter_state = filter_state_next
     x1, y1, x2, y2 = [int(coord[0]) for coord in (x1, y1, x2, y2)]
     x, y, w, h = x1, y1, x2-x1, y2-y1
     if args.debug:
         cv2.rectangle(
             debug_frame, (x, y), (x+w, y+h), (255, 0, 255), 2 * args.thickness
         )
 
     # Fix rectangle.
     x, y, w, h = pyffstream.fix_rect(args, x, y, w, h)
 
     # Determine resampling method.
     for i, (ratio, resample) in enumerate(args.resample_ratios):
         if min(w / args.output_width, h / args.output_height) >= ratio:
             break
     if args.debug:
         color_coeff = i / max(1, len(args.resample_ratios) - 1)
         color = (
             255 * (0 + color_coeff),
             255 * (1 - color_coeff),
             0,
         )
         cv2.rectangle(
             debug_frame, (x, y), (x+w, y+h), color, 2 * args.thickness,
         )
 
     # Cut and resize.
     output_frame = pyffstream.resize(
         frame[y:y+h, x:x+w], args.output_width, args.output_height, resample
     )
 
     # Return.
     return output_frame, debug_frame
 
 
 def main():
     pyffstream.run(__doc__, process, init, args_pre)
 
 
 if __name__ == '__main__':
     main()