#!/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.norm = None
		self.var = None
		self.final = None
		self.k_ref = 0.95
		self.k_var = 1
		self.k_norm = 1
		self.k_edge = 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 = 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, 5000, self.on_change_k_var)
		cv.createTrackbar('k_norm', 'matte', 0, 200, self.on_change_k_gain)
		cv.createTrackbar('k_edge', 'matte', 0, 200, self.on_change_k_edge)
		cv.createTrackbar('k_cc', 'matte', 0, 200, self.on_change_k_cc)
		cv.setMouseCallback('source', self.on_mouse_source)
		cv.setMouseCallback('final', self.on_mouse_final)
		self.update_controls()
		self.update_source()
		self.update_matte()
		self.update_edges()
		self.update_final()
		wait_key()
		print('Done')

	def on_mouse_source(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 on_mouse_final(self, event, x, y, flags, param):
		if event == cv.EVENT_LBUTTONDOWN:
			b = self.final[:, :, 0]*255
			g = self.final[:, :, 1]*255
			r = self.final[:, :, 2]*255
			a = self.matte*255
			bgra = cv.merge((b, g, r, a))
			cv.imwrite("image.png", bgra)
			print("Image written")

	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_norm', 'matte', int(self.k_norm * 100))
		cv.setTrackbarPos('k_edge', 'matte', int(self.k_edge * 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_norm = v / 100
		self.update_matte()

	def on_change_k_edge(self, v=0):
		self.k_edge = 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
		self.norm = np.clip(self.k_norm * np.linalg.norm(img_filt, axis=2), 0, 1)
		self.var = np.clip(self.k_var * np.var(img_filt, axis=2), 0, 1)
		self.update_edges()
		self.matte = np.clip(self.var + self.edges, 0, 1)
		self.win_matte.update(self.matte)
		self.update_final()

	def update_edges(self):
		self.edges = self.k_edge * self.norm * (1-self.var)
		self.win_edges.update(self.edges)

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


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