#!/usr/bin/env python


# Python 2/3 compatibility
from __future__ import print_function

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


def color_distance(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 despill(src, color_bg, edges):
	result = np.copy(src)

	result[:, :, 0] = src[:, :, 0] - edges * color_bg[0]
	result[:, :, 1] = src[:, :, 1] - edges * color_bg[1]
	result[:, :, 2] = src[:, :, 2] - edges * color_bg[2]

	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.edges = None
		self.img_final = None
		self.k_ref = 0.95
		self.k_var = 1
		self.k_gain = 1
		self.k_cc = 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.createTrackbar('k_cc', 'matte', 0, 100, self.on_change_k_cc)
		cv.setMouseCallback('source', self.on_mouse)
		self.update_controls()
		self.update_source()
		self.update_matte()
		self.update_edges()
		self.update_final()
		wait_key()
		print('Done')

	def on_mouse(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))
		cv.setTrackbarPos('k_cc', 'matte', int(self.k_cc * 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 on_change_k_cc(self, v=0):
		self.k_cc = 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.edges = self.matte * (1 - np.exp(-1 * (1 - self.matte)))
		self.win_edges.update(self.edges)
		self.update_final()

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


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