import os
import sys

import cv2 as cv
import numpy as np
import scipy.cluster
import scipy.ndimage
from matplotlib import pyplot as plt
import PIL.Image
from PIL.ImageDraw import ImageDraw

from annotations import DataFile,computeBoundingBox
from hough import show,houghLines


def createHistogram(img):
	# Convert BGR to HSV
	hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV)
	# H in range(0,180)
	# S in range(0,256)
	# V in range(0,256)
	planes=cv.split(hsv)
	hhist=cv.calcHist(planes,[0],None,[256],(0,180),accumulate=False)
	shist=cv.calcHist(planes,[1],None,[256],(0,256),accumulate=False)
	vhist=cv.calcHist(planes,[2],None,[256],(0,256),accumulate=False)

	width=512
	height=400
	binSize=width//256

	histImage = np.zeros((height, width, 3), dtype=np.uint8)
	cv.normalize(hhist, hhist, alpha=0, beta=height, norm_type=cv.NORM_MINMAX)
	cv.normalize(shist, shist, alpha=0, beta=height, norm_type=cv.NORM_MINMAX)
	cv.normalize(vhist, vhist, alpha=0, beta=height, norm_type=cv.NORM_MINMAX)

	for i in range(1, 256):
		cv.line(histImage, ( binSize*(i-1), height - int(round(hhist[i-1][0])) ),
						( binSize*(i), height - int(round(hhist[i][0])) ),
						( 255, 0, 0), thickness=2)
		cv.line(histImage, ( binSize*(i-1), height - int(round(shist[i-1][0])) ),
						( binSize*(i), height - int(round(shist[i][0])) ),
						( 0, 255, 0), thickness=2)
		cv.line(histImage, ( binSize*(i-1), height - int(round(vhist[i-1][0])) ),
						( binSize*(i), height - int(round(vhist[i][0])) ),
						( 0, 0, 255), thickness=2)

	cv.imshow('Source image', img)
	cv.imshow('calcHist Demo', histImage)
	cv.waitKey()


def quantize(img):
	arr=np.reshape(img,(-1,3)).astype(np.float)
	colors=np.array([[0,0,0],[255,255,255],[193,165,116]],np.float)
	print(colors)
	(centers,distortion)=scipy.cluster.vq.kmeans(arr,colors)
	print(centers)
	return centers


def computeClosest(x,centers):
	res=centers[0]
	d=np.linalg.norm(res-x)
	for c in centers:
		d_=np.linalg.norm(c-x)
		if d_<d:
			res=c
			d=d_
	return res


def score(arr1,arr2):
	try:
		return (arr1&arr2).sum() / ((arr1|arr2).sum() or 1)
	except TypeError:
		print(type(arr1),type(arr2))
		print(arr1.shape,arr2.shape)
		print(arr1.dtype,arr2.dtype)
		raise TypeError()

def maxOp55(arr):
	m=arr.max()
	return 1 if m>127 and arr[2,2]==m else 0

def ellipse(a,b):
	img=PIL.Image.new("1",(a,b))
	d=ImageDraw(img)
	d.ellipse((1,1,a-1,b-1),fill=1)
	img.save("/tmp/ellipse.png")
	return np.array(img,dtype=np.uint8)


filepath=sys.argv[1]
annotations=DataFile(sys.argv[2])
filename=os.path.basename(filepath)
(x1,y1,x2,y2)=computeBoundingBox(annotations[filename][0])
(w,h)=(x2-x1,y2-y1)
img=cv.imread(filepath)
(x3,x4,y3,y4)=(x1+w//4,x1+3*w//4,y1+h//4,y1+3*h//4)
print(x3,x4,y3,y4)
rect=img[y3:y4,x3:x4,:]
centers=quantize(rect)

for x in range(x1,x2):
	for y in range(y1,y2):
		pix=img[y,x]
		img[y,x]=computeClosest(pix,centers)

rect=img[y1:y2,x1:x2]
maskB=cv.inRange(rect,np.array([0,0,0]),np.array([80,80,80]))
maskB=cv.erode(maskB,np.ones((3,3),np.uint8))
# maskB=cv.erode(maskB,np.ones((3,3),np.uint8))
# maskB=cv.erode(maskB,np.ones((3,3),np.uint8))
maskW=cv.inRange(rect,np.array([160,160,160]),np.array([256,256,256]))
maskW=cv.erode(maskW,np.ones((3,3),np.uint8))
# maskW=cv.erode(maskW,np.ones((3,3),np.uint8))
# maskW=cv.erode(maskW,np.ones((3,3),np.uint8))

show(img,filename)
show(maskB,filename)
show(maskW,filename)
stones=cv.bitwise_or(maskB,maskW)
# houghLines(stones)

(bh,bw)=stones.shape
sw=bw//19
sh=bh//19
print(stones.shape,(sw,sh))
ell=ellipse(sw,sh)*255
# print(ell)
hitMap=np.zeros_like(stones,dtype=np.uint8)
for i in range(sw,bw):
	for j in range(sh,bh):
		region=stones[j-sh:j, i-sw:i]
		hitMap[j,i]=255*score(region,ell)
show(hitMap)

gridMap=np.zeros_like(hitMap,dtype=np.uint8)
for i in range(5,bw):
	for j in range(5,bh):
		region=hitMap[j-5:j, i-5:i]
		gridMap[j,i]=255*maxOp55(region)
show(gridMap)

houghLines(gridMap)
# ministones=cv.resize(stones,None,fx=0.25,fy=0.25,interpolation=cv.INTER_AREA)
# dft = cv.dft(np.float32(ministones),flags = cv.DFT_COMPLEX_OUTPUT)
# dft_shift = np.fft.fftshift(dft)
# magnitude_spectrum = 20*np.log(cv.magnitude(dft_shift[:,:,0],dft_shift[:,:,1]))
# plt.subplot(121),plt.imshow(stones, cmap = 'gray')
# plt.title('Input Image'), plt.xticks([]), plt.yticks([])
# plt.subplot(122),plt.imshow(magnitude_spectrum, cmap = 'gray')
# plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
# plt.show()