#!/usr/bin/env python
from __future__ import print_function
import sys, argparse, signal
import cv2
import cv2.cv as cv
import numpy as np

screen_width = 1920
screen_height = 1080
pattern_size = 4
pattern_mode = 1

#import subprocess

def nothing(par):
    return
    
def signal_handler(signal, frame):
    sys.exit(0)
    
def draw_pattern():
    global pattern_size
    if pattern_size <= 1:
        pattern_size = 1
    tmp_img = np.zeros((screen_height,screen_width,3), np.uint8)
    for j in range(0,screen_height,2*pattern_size):
        if (pattern_mode ==1):
            for i in range(0,screen_width,2*pattern_size):
                cv2.rectangle(tmp_img,(i,j),(i+pattern_size,j+pattern_size),(255,255,255),cv.CV_FILLED)
                cv2.rectangle(tmp_img,(i+pattern_size,j+pattern_size),(i+2*pattern_size,j+2*pattern_size),(255,255,255),cv.CV_FILLED)
        if (pattern_mode ==2):
                cv2.rectangle(tmp_img,(0,j),(screen_width,j+pattern_size),(255,255,255),cv.CV_FILLED)
    return tmp_img

if __name__ == '__main__':
    # register signal handler for a clean exit    
    signal.signal(signal.SIGINT, signal_handler)
    
    # command line parser
    parser = argparse.ArgumentParser(description='Python Schlieren Imager')
    parser.add_argument('-i', '--input_source', default=0, help='Input Source, either filename or camera index, default = 0')
    parser.add_argument('-iw', '--input_width', default=0, help='Width of captured frames')
    parser.add_argument('-ih', '--input_height', default=0, help='Height of captured frames')
    args = parser.parse_args()

    # open video and get cascade ready
    if str(args.input_source).isdigit():
        video_src = int(args.input_source)
    else:
        video_src = args.input_source
    # uncomment following line and set to default device for testing.
    video_src = 1
    cam = cv2.VideoCapture(video_src)  
    if args.input_width > 0 and args.input_height > 0:
        cam.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, float(args.input_width))
        cam.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, float(args.input_height))
    if cam is None or not cam.isOpened():
        print("Warning: unable to open video source:" + video_src, file=sys.stderr)
    frame_width = int(cam.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH))
    ret, img = cam.read()
    if ret == False:
        print("No more frames or capture device down - exiting.", file=sys.stderr)
        sys.exit(0)
    vis = img.copy()
    ref = cv2.cvtColor(vis, cv2.COLOR_BGR2GRAY)
            
    # create window and GUI
    cv2.namedWindow('Python Schlieren Imager', 0)
    cv2.namedWindow('Raw Image', 0)
    cv2.namedWindow('Background', 0)
    
    # create empty image for background pattern
    pattern = np.zeros((screen_height,screen_width,3), np.uint8)
    
    # main video processing loop
    while True:
        ret, img = cam.read()
        if ret == False:
            print("No more frames or capture device down - exiting.", file=sys.stderr)
            sys.exit(0)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        gray = cv2.equalizeHist(gray)
        t = cv2.getTickCount() / cv2.getTickFrequency()      
        vis = cv2.absdiff(gray, ref)
        vis = cv2.equalizeHist(vis)
        vis = cv2.applyColorMap(vis, cv2.COLORMAP_RAINBOW)
        cv2.imshow('Python Schlieren Imager', vis)
        cv2.imshow('Raw Image', img)
        cv2.imshow('Background', pattern)
        keyscan = 0xFF & cv2.waitKey(5)
        if keyscan == 113:
            break
        if keyscan == 99:
            ref = gray.copy()
        if keyscan == 49: # 1 chequerboard pattern
            pattern_mode = 1
            pattern = draw_pattern()
        if keyscan == 50: # 2 striped pattern
            pattern_mode = 2
            pattern = draw_pattern()
        if keyscan == 61: # + increase size
            pattern_size += 1
            pattern = draw_pattern()
        if keyscan == 45: # - decrease size
            pattern_size -= 1
            pattern = draw_pattern()
    sys.exit(0)