import os
import sys
import random

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("k-means centers:",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)

def detectStones(img):
	(bh,bw)=img.shape
	sw=bw//19
	sh=bh//19
	print(img.shape,(sw,sh))
	ell=ellipse(sw,sh)*255
	# print(ell)
	hitMap=np.zeros_like(img,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)
	return hitMap

def detectGrid(img):
	(bh,bw)=img.shape

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

def locateStone(img):
	(bh,bw)=img.shape
	sw=bw//19
	sh=bh//19
	print(img.shape,(sw,sh))
	ell=ellipse(sw,sh)*255
	# print(ell)

	y=random.randrange(sh,bh)
	x=random.randrange(sw,bw)
	region=stones[y-sh:y, x-sw:x]
	sc=score(region,ell)
	print(sc)
	show(region)
	return sc

def locateLocalMaximum(img):
	(bh,bw)=img.shape
	x=random.randrange(0,bw)
	y=random.randrange(0,bh)
	val=img[y,x]
	img_=cv.cvtColor(img,cv.COLOR_GRAY2BGR)
	while True:
		for (dx,dy) in [(0,1),(1,0),(0,-1),(-1,0),(1,1),(-1,1),(1,-1),(-1,-1)]:
			x_=x+dx
			y_=y+dy
			val_=img[y_,x_] if 0<=x_<bw and 0<=y_<bh else 0
			if val_>val:
				x=x_
				y=y_
				val=val_
				img_.itemset((y,x,0),255)
				continue
		break
	print(x,y,val)
	img_.itemset((y,x,2),255)
	show(img_)
	return (x,y,val)


if __name__=="__main__":
	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:",x3,x4,y3,y4)
	rect=img[y3:y4,x3:x4,:]
	centers=quantize(rect)
	print("x1,x2,y1,y2:",(x1,x2,y1,y2))
	data=np.reshape(img[y1:y2,x1:x2,:],(-1,3))
	print("data.shape:",data.shape)
	(keys,dists)=scipy.cluster.vq.vq(data,centers)
	print("keys.shape:",keys.shape)
	pixels=np.array([centers[k] for k in keys],dtype=np.uint8).reshape((y2-y1,x2-x1,3))
	img[y1:y2,x1:x2,:]=pixels
	print("image quantized")

	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),iterations=2)
	maskW=cv.inRange(rect,np.array([160,160,160]),np.array([256,256,256]))
	maskW=cv.erode(maskW,np.ones((3,3),np.uint8),iterations=2)

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

	(contours,hierarchy)=cv.findContours(stones,cv.RETR_LIST,cv.CHAIN_APPROX_SIMPLE)
	contourImg=cv.drawContours(cv.cvtColor(stones,cv.COLOR_GRAY2BGR), contours, -1, (0,255,0), 1)
	for (i,c) in enumerate(contours):
		print(i)
		moments=cv.moments(c)
		center=(moments["m10"]/(moments["m00"] or 1), moments["m01"]/(moments["m00"] or 1))
		print("center:", center)
		cv.circle(contourImg,tuple(map(int,center)),3,(255,255,0))
		area=cv.contourArea(c)
		print("area:",area)
		(x,y,w,h)=cv.boundingRect(c)
		print("bounding box:",(x,y,w,h))
		print("coverage1:",area/(w*h or 1))
		hull=cv.convexHull(c)
		print("coverage2:",area/(cv.contourArea(hull) or 1))
		print()
	show(contourImg)

	# for i in range(10):
	# 	locateStone(stones)
	# distMap=cv.distanceTransform(stones,cv.DIST_L2,5)
	# print("dist map:")
	# show(distMap)
	# print("hit map:")
	# hitMap=detectStones(stones)
	# for i in range(10):
	# 	locateLocalMaximum(hitMap)
# dlib.find_peaks

# 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()