from pickletools import uint8

import cv2 as cv
import numpy as np
from matplotlib import pyplot as plt

def key_ish(kid: int, color: np.array):
	prg = [(1,2), (0,2), (0,1)]
	res = 0

	c = prg[kid]
	ci1 = c[0]
	ci2 = c[1]
	c1 = color[ci1]
	c2 = color[ci2]
	k0_1 = color[kid]
	k0_2 = color[kid]
	k0_12 = color[kid]
	s12 = c1 + c2
	if c1 > 0:
		k0_1 = color[kid] / c1
	if c2 > 0:
		k0_2 = color[kid] / c2
	if s12 > 0:
		k0_12 = color[kid] / s12

#	if s12 > 2*c1 and s12 > 2*c2:
#		res = k0_12

	if c1 > c2:
		res = k0_1
	else:
		res = k0_2

	return res

def despill(src, key_color, mat: float):
	result = np.copy(src)

	result[:, :, 0] -= (1-mat)*key_color[:, :, 0]
	result[:, :, 1] -= (1-mat)*key_color[:, :, 1]
	result[:, :, 2] -= (1-mat)*key_color[:, :, 2]

	return result

def main():
	key_index = 1
	key_bgr = np.float32([[[0, 0, 0]]])
	key_bgr[:,:,key_index] = 1

	orig = np.ones([200, 360, 3], dtype=np.float32)
	for x, hn in enumerate(np.linspace(0, 1, 360)):
		for y in range(200):
			s = y/200
			orig_hsv = np.float32([[[hn * 360, s, 1]]])
			orig_bgr = cv.cvtColor(orig_hsv, cv.COLOR_HSV2BGR)
			orig[y, x] = orig_bgr

	cv.imshow("orig", orig)

	k_tol = 0.3
	k_boost = 1.5
	_x = []
	_y = []
	pic = np.ones([200, 360, 3], dtype=np.float32)
	mat = np.ones([200, 360, 1], dtype=np.float32)
	for x, hn in enumerate(np.linspace(0, 1, 360)):
		_x.append(x)
		v = 0
		for y in range(200):
			s = y/200
			spec_hsv = np.float32([[[hn * 360, s, 1]]])
			spec_bgr =  cv.cvtColor(np.array(spec_hsv, dtype=np.float32),cv.COLOR_HSV2BGR)
			v0 = spec_bgr[0,0,:]
			v1 = min(1.0, k_boost * (1.0 - min(1.0, k_tol*key_ish(key_index, v0))))
			pic_hsv = np.float32([[[hn*360,s,v1]]])
			pic_bgr = cv.cvtColor(pic_hsv,cv.COLOR_HSV2BGR)
			pic_bgr = despill(pic_bgr, key_bgr, v1)
			pic[y, x] = pic_bgr
			mat[y, x] = v1
		_y.append(v1)
	print(mat)
	cv.imshow("pic", pic)
	cv.imshow("mat", mat)

	plt.plot(_x, _y)
	plt.grid()
#	plt.show()

if __name__ == '__main__':
	main()
	ch = cv.waitKey()