# import the necessary packages
from __future__ import print_function
from PiVideoStream import PiVideoStream
from imutils.video import FPS
from picamera.array import PiRGBArray
from picamera import PiCamera
import numpy as np
import argparse
import imutils
import time
import cv2
import pprint
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

width = 320
height = 240

class FindObjects():
	def __init__(self):
		self.innerOuterRatio = 0.6	# const
		self.meanThickness = 0
		self.meanOuterDiameter = 0
		self.meanInnerDiameter = 0
		self.maxDiameter = -10000
		self.minDiameter = +10000
		self.pp = pprint.PrettyPrinter(indent=4)

	@staticmethod
	def getFamily(family, parentId, hierachy):
	
		family.append(parentId)
		if hierachy[0][parentId][2] != -1:
			return FindObjects.getFamily(family, hierachy[0][parentId][2], hierachy)
		else:
			return family


	def find(self, contours, hierachy):
		count = 0
		objects = []
		temp_objects = []
		objectsIds = []
		objectsCandIds = []
		for cnt in contours:
			if hierachy[0][count][3] == -1:
				family = FindObjects.getFamily([], count, hierachy)

				dupDict = {}
				members = []
				for member in family:
					x,y,w,h = cv2.boundingRect(contours[member]);
					key = str([x,y,w,h]) + '.key'				
					if not key in dupDict:
						dupDict[key] = member
						members.append({ 'id' : member, 'bbox' : [x,y,w,h]})

				objectsCandIds.append(members)

			count += 1
	
		for family in objectsCandIds:
			obj = []
			for member in family:
				area = cv2.contourArea(contours[member['id']])
				perimeter = cv2.arcLength(contours[member['id']], True)
				pi = 3.14159265359
				Q = 0
				if (perimeter > 0):
					Q = 4*pi*area/(perimeter*perimeter)

				if Q >= 0.7:
					diameter = max(member['bbox'][2], member['bbox'][3])
					self.maxDiameter = max(self.maxDiameter, diameter)
					self.minDiameter = min(self.minDiameter, diameter)
					member['diameter'] = diameter
					member['pos'] = (int(member['bbox'][0] + member['bbox'][2]/2), int(member['bbox'][1] + member['bbox'][3]/2))
					obj.append(member)

			if obj:
				temp_objects.append(obj)

		for obj in temp_objects:
			for member in obj:
				if member['diameter'] < self.innerOuterRatio*self.maxDiameter:
					member['isHole'] = True
				else:
					member['isHole'] = False

			if len(obj) == 2:
				self.meanThickness = int(abs(obj[0]['diameter'] - obj[1]['diameter'])/2)

			objects.append({'thickness' : self.meanThickness, 'members' : obj})
						
		return objects

	def printStats(self):
		print ("maxDiameter   = " + str(self.maxDiameter) + " px")
		print ("minDiameter   = " + str(self.minDiameter) + " px")
		print ("meanThickness = " + str(self.meanThickness) + " px")
		
# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-n", "--num-frames", type=int, default=100,
	help="# of frames to loop over for FPS test")
ap.add_argument("-d", "--display", type=int, default=-1,
	help="Whether or not frames should be displayed")
args = vars(ap.parse_args())

# created a *threaded *video stream, allow the camera sensor to warmup,
# and start the FPS counter
print("[INFO] sampling THREADED frames from `picamera` module...")
vs = PiVideoStream(resolution=(width,height), framerate=30).start()
time.sleep(2.0)
fps = FPS().start()
findObjects = FindObjects()

beadColors = []

# loop over some frames...this time using the threaded stream
while fps._numFrames < args["num_frames"]:
	# grab the frame from the threaded video stream
	frame = vs.read()
	gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
	gray_blurred = cv2.medianBlur(gray,5)

	# Canny edge detection
	img1_canny = cv2.Canny(gray_blurred, 100, 50)
	kernel = np.ones((1,1),np.uint8)
#	img1_canny = cv2.blur(img1_canny,(3,3))
#	img1_canny = cv2.dilate(img1_canny,kernel,iterations = 1)

	# Thresholding
#	ret,img1_thr = cv2.threshold(gray_blurred,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
#	img1_thr = cv2.adaptiveThreshold(gray_blurred,255,cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY,11,2)
#	img1_thr = cv2.adaptiveThreshold(gray_blurred,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2)

	# Contours
	(_, contours, hierachy) = cv2.findContours(img1_canny.copy(),cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

	objects = findObjects.find(contours, hierachy)
	
#	print (ids)

	img1_objects = frame.copy()
	img1_colors = np.zeros((height,width,3), np.uint8)
	maskCenter = 4
	roides = []
	beads = []
	for obj in objects:
		memberCount = 0
		radius = 0
		roi = 0
		for member in obj['members']:
			# Draw bounding box
			x,y,w,h = [member['bbox'][0], member['bbox'][1], member['bbox'][2], member['bbox'][3]];
			if member['isHole']:
				img1_objects = cv2.rectangle(img1_objects,(x,y),(x+w,y+h),(255,0,0),2)
				thickness = obj['thickness']-maskCenter
				diameter = member['diameter']+thickness+maskCenter
				roi = frame[y-thickness:y+diameter, x-thickness:x+diameter]
				pos = member['pos']
				if diameter > 0:
					mask = np.zeros((diameter,diameter,1), np.uint8)
					mask = cv2.circle(mask,(int(diameter/2), int(diameter/2)),radius,255,thickness)
					mean_color = cv2.mean(roi)
					img1_colors = cv2.circle(img1_colors,(int(pos[0]),int(pos[1])),int(diameter/2),mean_color,-1)
					beads.append({'pos' : pos, 'diameter' : diameter, 'color' : mean_color[0:3]})
					img1_objects = cv2.circle(img1_objects,member['pos'],int(diameter/2),(255,255,255),thickness)
			else:
				img1_objects = cv2.rectangle(img1_objects,(x,y),(x+w,y+h),(0,0,255),2)

		memberCount += 1

	for bead in beads:
		beadColors.append(bead['color'])

#	cv2.imshow('Thresholded',img1_thr)
	cv2.imshow('detected colors', img1_colors)
	cv2.imshow('Canny',img1_canny)
	cv2.imshow('Objects',img1_objects)
	cv2.waitKey(1)

   	# update the FPS counter
	fps.update()

# stop the timer and display FPS information
fps.stop()
vs.stop()
print("[INFO] elasped time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
print("[INFO] Camera : elasped time: {:.2f}".format(vs.getfps().elapsed()))
print("[INFO] Camera: approx. FPS: {:.2f}".format(vs.getfps().fps()))

# Output stats
findObjects.printStats()

# Scatter plot of colors
numObjects = len(beadColors);
blue = np.zeros(numObjects)
green = np.zeros(numObjects)
red = np.zeros(numObjects)
plotColors = np.zeros((numObjects,3))
i = 0
for color in beadColors:
	blue[i] = color[0]/255
	green[i] = color[1]/255
	red[i] = color[2]/255
	plotColors[i] = [red[i], green[i], blue[i]]
	i += 1

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(blue, green, red, zdir='z', s=50, facecolors=plotColors, lw = 0, depthshade=True)
plt.show()

# do a bit of cleanup
cv2.destroyAllWindows()