#!/usr/bin/env python
'''
Chroma keying in OpenCV

'''
# Python 2/3 compatibility
from __future__ import print_function

import numpy as np
import sys
import argparse
import cv2 as cv


def colordistance(color1, color2):
	d0 = (color1[0] - color2[0])
	d1 = (color1[1] - color2[1])
	d2 = (color1[2] - color2[2])
	return np.math.sqrt(d0 * d0 + d1 * d1 + d2 * d2)

def blend(src, matte):
	result = np.copy(src)

	result[:, :, 0] = src[:, :, 0] * matte
	result[:, :, 1] = src[:, :, 1] * matte
	result[:, :, 2] = src[:, :, 2] * matte

	return result

def wait_key():
	while True:
		ch = cv.waitKey()
		if ch == 27:
			break

class ImageWindow():
	def __init__(self, name):
		self.name = name
		cv.namedWindow(self.name)

	def update(self, img):
		cv.imshow(self.name, img)

class Chromakey():
	def __init__(self):
		self.matte = None
		self.k_ref = 0.95
		self.k_var = 1
		self.k_gain = 1
		self.color_key = np.array([0.9, 0.2, 0.1])
		ap = argparse.ArgumentParser()
		ap.add_argument("-f", "--filename", type=str, default=None,
						help="Input image")

		args = vars(ap.parse_args())
		inputfilename = args["filename"]

		self.img_source = np.float32(cv.imread(inputfilename)) / 255
		if self.img_source is None:
			print('Failed to load file:', inputfilename)
			sys.exit(1)

		self.win_source = ImageWindow('source')
		self.win_matte = ImageWindow('matte')
		self.win_final = ImageWindow('final')
		self.win_edges = ImageWindow('edges')
		cv.createTrackbar('k_ref', 'matte', 0, 150, self.on_change_k_ref)
		cv.createTrackbar('k_var', 'matte', 0, 1000, self.on_change_k_var)
		cv.createTrackbar('k_gain', 'matte', 0, 2000, self.on_change_k_gain)
		cv.setMouseCallback('source', self.onMouse)
		self.update_controls()
		self.update_source()
		self.update_matte()
		self.update_edges()
		self.update_final()
		wait_key()
		print('Done')

	def onMouse(self, event, x, y, flags, param):
		if (event == cv.EVENT_LBUTTONDOWN):
			self.color_key = self.img_source[y,x]
			print("color_key", self.color_key)
			self.update_matte()

	def update_controls(self):
		cv.setTrackbarPos('k_ref', 'matte', int(self.k_ref*100))
		cv.setTrackbarPos('k_var', 'matte', int(self.k_var*100))
		cv.setTrackbarPos('k_gain', 'matte', int(self.k_gain*100))

	def on_change_k_ref(self, v=0):
		self.k_ref = v/100
		self.update_matte()

	def on_change_k_var(self, v=0):
		self.k_var = v/100
		self.update_matte()

	def on_change_k_gain(self, v=0):
		self.k_gain = v/100
		self.update_matte()

	def update_source(self):
		self.win_source.update(self.img_source)

	def update_matte(self):
		img_filt = self.img_source - self.k_ref * self.color_key
		img_norm = np.linalg.norm(img_filt, axis=2)
		img_var = np.var(img_filt, axis=2)
		img_exp = self.k_gain*(1-np.exp(-self.k_var*img_var))
		self.matte = np.clip(img_exp, 0, 1)
		self.win_matte.update(self.matte)
		self.update_edges()

	def update_edges(self):
		img_filt = self.img_source - self.k_ref * self.color_key
		img_norm = np.linalg.norm(img_filt, axis=2)
		edges = np.clip(img_norm-self.matte, 0, 1)
		self.win_edges.update(edges)
		self.update_final()

	def update_final(self):
		# Multiply the foreground with the alpha matte
		self.img_final = blend(self.img_source, self.matte)
		self.win_final.update(self.img_final)

if __name__ == '__main__':
	print(__doc__)
	ck = Chromakey()
	cv.destroyAllWindows()