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

#scikit-video (scikit-video.org) is used for recording because of OpenCV Linux bug
from skvideo.io import FFmpegWriter as VideoWriter

screen_width = 1920
screen_height = 1080
pattern_size = 4
pattern_mode = 2
color_mode = -1

#switches for image filters and tools
showhud = True
refdiff = True
display = False
histgrm = False
pseudoc = False
avgblur = False
denoise = False
recordi = False
recordv = False
novfile = True
imgindx = 0

#more constants
denoise_strength = 33
stacklen = 10
hud_src = ''
hud_filter = ''
hud_rec = ''
hud_color = ''
image_dst = ''
video_dst = 'output.avi'
     
def nothing(par):
    return
    
def signal_handler(signal, frame):
    sys.exit(0)
    
def img_stack(imgsrc):
    for i in range(0,stacklen):
        ret, tmpimg = imgsrc.read()
        if ret == False:
            break
        tmpvis = tmpimg.copy()
        if i == 0:
            hdvis = np.float32(tmpvis)
        else:
            hdvis = hdvis + np.float32(tmpvis)
    stacked = np.uint8(hdvis/stacklen)
    stacked = cv2.cvtColor(stacked, cv2.COLOR_BGR2GRAY)
    stacked = clahe.apply(stacked)
    #stacked = cv2.equalizeHist(stacked)
    return stacked
    
def draw_pattern():
    global pattern_size
    if pattern_size <= 1:
        pattern_size = 1
    tmp_img = np.zeros((screen_height,screen_width,3), np.uint8)
    if (pattern_mode == 3):
        for j in range(0,screen_width,2*pattern_size):
            cv2.rectangle(tmp_img,(j,0),(j+pattern_size,screen_height),(255,255,255),cv.CV_FILLED)
    else:
        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
    
def update_hudtext():
    global hud_src
    global hud_filter
    global hud_rec
    global hud_color
    if display:
        hud_src = '(D)isplay: Rendered Output'
        hud_color = '(C)olor Mode: '
        if pseudoc == False:
            hud_color += 'Greyscale'
        else:
            hud_color += str(color_mode)
        hud_filter = 'Filters (E,B,N): '
        if histgrm:
            hud_filter += '(E)qualize '
        if avgblur:
            hud_filter += '(B)lur '
        if denoise:
            hud_filter += '(N)oise '
    else:
        hud_src = '(D)isplay: Raw Input'
        hud_color = ''
        hud_filter = ''
    if recordv:
        hud_rec = 'Recording (V)ideo! '
    else:
        hud_rec = ''
    if recordi:
        hud_rec += 'Recording (I)mages!'

if __name__ == '__main__':
    # color definitions for text in HUD
    green = (0, 255, 0)
    red = (0, 0, 255)
    
    # 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=1024, help='Width of captured frames')
    parser.add_argument('-ih', '--input_height', default=768, help='Height of captured frames')
    parser.add_argument('-ov', '--output_video', default='none', help='Save output as video (Path or Filename).')
    parser.add_argument('-oi', '--output_images', default='none', help='Save output as image sequence (Path or Prefix).')
    parser.add_argument('-fe', '--filter_equalize', action='store_true', help='Enable filter: Histrogram Equalization')    
    parser.add_argument('-fb', '--filter_blur', action='store_true', help='Enable filter: Blur')
    parser.add_argument('-fd', '--filter_denoise', default='none', help='Enable Denoising (filter strength)')
    parser.add_argument('-dd', '--disable_delta', action='store_true', help='Disable Schlieren Differentiation')    
    parser.add_argument('-c', '--color_mode', default=color_mode, help='Pseudocolor mode (1-11)')
    parser.add_argument('-pm', '--pattern_mode', default=pattern_mode, help='Background Pattern (1-3)')
    parser.add_argument('-ps', '--pattern_size', default=pattern_size, help='Background Pattern Size')
    args = parser.parse_args()

    # process command line arguments
    if args.output_video != 'none':
        recordv = True
        video_dst = args.output_video
        
    if args.output_images != 'none':
        recordi = True
        image_dst = args.output_images
    
    histgrm = args.filter_equalize
    avgblur = args.filter_blur
    refdiff = not args.disable_delta
    
    if args.filter_denoise != 'none':
        denoise = True
        denoise_strength = float(args.filter_denoise)
    
    if int(args.color_mode) <= 11 and int(args.color_mode) >= 0:
        color_mode = int(args.color_mode)
        pseudoc = True
            
    pattern_mode = int(args.pattern_mode)
        
    if int(args.pattern_size) > 0:
        pattern_size = int(args.pattern_size)
    
    # open video
    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))
    frame_height = int(cam.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT))
    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()
            
    # we wait a second to make sure whoever started this program has released RETURN
    time.sleep(1)
    # get value for "no key pressed"
    nokey = cv2.waitKey(5)
    
    # create window and GUI
    cv2.namedWindow('Python Schlieren Imager', 0)
    cv2.namedWindow('Background', 0)
    update_hudtext()
    
    # create image for background pattern
    pattern = np.zeros((screen_height,screen_width,3), np.uint8)
    pattern = draw_pattern()
    cv2.imshow('Background', pattern)
    
    # create a CLAHE object (Contrast Limited Adaptive Histogram Equalization)
    clahe = cv2.createCLAHE(clipLimit=4.0, tileGridSize=(4,4))
    
    # use first frames for reference image.
    if refdiff:
        ref = img_stack(cam)
    else:
        ref = ''
    
    # 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 = clahe.apply(gray)        
        if refdiff:
            vis = cv2.absdiff(gray, ref)
        else:
            vis = gray
        if histgrm:
            vis = cv2.equalizeHist(vis)
        if denoise:
            cv2.fastNlMeansDenoising(vis, vis, denoise_strength, 7, 21)
        if avgblur:
            vis = cv2.blur(vis,(5,5))
        if pseudoc:
            vis = cv2.applyColorMap(vis, color_mode)
        else:
            vis = cv2.cvtColor(vis, cv2.COLOR_GRAY2BGR)
        if recordv:
            if novfile:
                video = VideoWriter(video_dst)
                novfile = False
            video.writeFrame(vis)
        if recordi:
            filename = image_dst + str(imgindx).zfill(8) + '.bmp'
            cv2.imwrite(filename, vis)
            imgindx += 1
        if display:
            dsp = vis.copy()
        else:
            dsp = img.copy()
        if showhud:
            cv2.putText(dsp, hud_src, (20, 50), cv2.FONT_HERSHEY_PLAIN, 2.0, green, thickness=2, lineType=cv2.CV_AA)
            cv2.putText(dsp, hud_filter, (20, 80), cv2.FONT_HERSHEY_PLAIN, 2.0, green, thickness=2, lineType=cv2.CV_AA)
            cv2.putText(dsp, hud_color, (20, 110), cv2.FONT_HERSHEY_PLAIN, 2.0, green, thickness=2, lineType=cv2.CV_AA)
            cv2.putText(dsp, hud_rec, (20, 140), cv2.FONT_HERSHEY_PLAIN, 2.0, red, thickness=2, lineType=cv2.CV_AA)        
        cv2.imshow('Python Schlieren Imager', dsp)        
        keyscan = cv2.waitKey(5)
        if keyscan != nokey:
            asckey = 0xFF & keyscan
            if asckey == 113: # q terminates program
                break
            if asckey == 114: # r capture reference image
                ref = img_stack(cam)
                histgram = np.amax(ref)
                #ref = gray.copy()
            if asckey == 112: # p change background pattern
                pattern_mode += 1
                if pattern_mode > 3:
                    pattern_mode = 1
                pattern = draw_pattern()
            if asckey == 43: # + increase size
                pattern_size += 1
                pattern = draw_pattern()
            if asckey == 45: # - decrease size
                pattern_size -= 1
                pattern = draw_pattern()
            if asckey == 101: # e toggle histogram equalizer
                histgrm = not histgrm
            if asckey == 99: # c choose color palette
                color_mode += 1
                if color_mode > 11:
                    pseudoc = False
                    color_mode = -1
                else:
                    pseudoc = True
            if asckey == 98: # b toggle median blur
                avgblur = not avgblur
            if asckey == 118: # v toggle video recording
                recordv = not recordv
            if asckey == 105: # i toggle image sequence recording
                recordi = not recordi
            if asckey == 104: # h toggle HUD
                showhud = not showhud
            if asckey == 100: # d toggle display
                display = not display
            if asckey == 115: # s toggle schlieren differentiation
                refdiff = not refdiff
            if asckey == 110: # n toggle denoise filter
                denoise = not denoise
            cv2.imshow('Background', pattern)
            update_hudtext()
    if novfile == False:
        video.close()
    sys.exit(0)