# import the necessary packages
from __future__ import print_function

from imutils.video import FPS

import sys
import numpy as np
import numpy.matlib
import argparse
import cv2
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

width = 320
height = 240
		
# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-n", "--num-frames", type=int, default=100,
	help="# of frames to loop over for FPS test")
ap.add_argument("-r", "--framerate", type=int, default=30,
	help="# of frames to loop over for FPS test")
ap.add_argument("-f", "--filename", type=str, default='piCamera',
	help="Whether or not frames should be displayed")
ap.add_argument("-c", "--calibrate", type=str, default='off',
	help="Whether calibration shall be performed")
ap.add_argument("-u", "--use-calibration", type=str, default='off',
	help="Whether calibration shall be used")
args = vars(ap.parse_args())

videoFile = args["filename"]
framerate = args["framerate"]

# created a *threaded *video stream, allow the camera sensor to warmup,

frame = np.float32(cv2.imread(videoFile))/255
h, w, nc = frame.shape

if args["calibrate"] == 'on':
	lens_corr = np.ones((h, w, nc), np.float32)
	
	roi_f32 = frame[int(height/2-8):int(height/2+8), int(width/2-8):int(width/2+8), :]
	
	ref_color = cv2.mean(roi_f32)
	ref_rr_rg_rb = (1.0, ref_color[1]/ref_color[2], ref_color[0]/ref_color[2], 1.0)
	ref_rr_rg_rb = (1.0, 1.0, 1.0, 1.0)
	red = frame[:,:,2]
	red_div = np.reshape(np.matlib.repmat(red, 1, nc), (h,w,nc), order='F')
	
	rr_rg_rb = cv2.divide(red_div, frame.copy())
	lens_corr0 = cv2.divide(rr_rg_rb, ref_rr_rg_rb)
	lens_corr = cv2.divide(lens_corr0, red_div)/1.9
	lens_test = cv2.multiply(frame.copy(), lens_corr)
	params = list()
	params.append(cv2.IMWRITE_PNG_COMPRESSION)
	params.append(0)

	cv2.imwrite("rr_rg_rb.png", np.uint8(rr_rg_rb*128.0), params)
	cv2.imwrite("lens_corr0.png", np.uint8(lens_corr0*128.0), params)
	cv2.imwrite("lens_corr.png", np.uint8(lens_corr*128.0), params)
	cv2.imwrite("lens_test.png", np.uint8(lens_test*255.0), params)
	cv2.imwrite("red_div.png", np.uint8(red_div*255.0), params)

else:
	lens_corr = np.float32(cv2.imread("lens_corr.png"))/128
	lens_test = cv2.multiply(frame, lens_corr)
	params = list()
	params.append(cv2.IMWRITE_PNG_COMPRESSION)
	params.append(0)

	cv2.imwrite("lens_test.png", np.uint8(lens_test*255.0), params)
	
# do a bit of cleanup
cv2.destroyAllWindows()