#!/usr/bin/python3
from PyQt5.QtWidgets import QApplication, QWidget, QShortcut
from PyQt5.QtGui import QKeySequence
from PyQt5.QtCore import Qt, QThread, QPoint
import time
import cv2
from picamera2 import MappedArray, Picamera2, Preview
from picamera2.previews.qt import QGlPicamera2
import numpy as np
from pythonosc import udp_client


def rectFromPoint(center, len, width, axis):
	rect = ((0, 0), (0, 0))
	l = int(len/2)
	w = int(width/2)
	if(axis == 'x'):
		rect = ((center[0] - l, center[1] - w), (center[0] + l, center[1] + w))
	elif(axis == 'y'):
		rect = ((center[0] - w, center[1] - l), (center[0] + w, center[1] + l))
	return rect


def rectsFromPoint(center, l1, l2, l3, w, axis):
	centerFine = center
	fineInner = rectFromPoint(centerFine, l1, w, axis)
	if(axis == 'x'):
		fineInnerNeg = rectFromPoint((centerFine[0] - int(l1 / 4), centerFine[1]), int(l1 / 2), w, axis)
		fineInnerPos = rectFromPoint((centerFine[0] + int(l1 / 4), centerFine[1]), int(l1 / 2), w, axis)
	elif(axis == 'y'):
		fineInnerNeg = rectFromPoint((centerFine[0], centerFine[1] - int(l1 / 4)), int(l1 / 2), w, axis)
		fineInnerPos = rectFromPoint((centerFine[0], centerFine[1] + int(l1 / 4)), int(l1 / 2), w, axis)

	fineOuter = rectFromPoint(centerFine, l2, w, axis)
	
	centerCoarse = center
	if(axis == 'x'):
		centerCoarse = (center[0], center[1] + w)
	elif(axis == 'y'):
		centerCoarse = (center[0] + w, center[1])
	coarse = rectFromPoint(centerCoarse, l3, w, axis)
		
	return [fineInnerNeg, fineInnerPos, fineInner, fineOuter, coarse, center]
	
	
def moveROI(event, x, y, flags, params):
	global roiX, roiY, moving, l1, l2, l3, w, selectedAxis
	if event == cv2.EVENT_LBUTTONDOWN:
		moving = True

	elif event==cv2.EVENT_MOUSEMOVE:
		if moving==True:
			if(selectedAxis == 'x'):
				roiX = rectsFromPoint((x, y), l1, l2, l3, w, selectedAxis)
			elif(selectedAxis == 'y'):
				roiY = rectsFromPoint((x, y), l1, l2, l3, w, selectedAxis)

	elif event == cv2.EVENT_LBUTTONUP:
		moving = False


def crop(frame, rect):
	return frame[rect[0][1]:rect[1][1], rect[0][0]:rect[1][0]]


def replaceCrop(frame, rect, crop):
	frame[rect[0][1]:rect[1][1], rect[0][0]:rect[1][0]] = crop


def genDKernel(dVal):
	return np.ones((dVal, dVal), np.uint8)


def track(frame, roi, dKernel):
	exp = 2
	roiFineInnerNeg, roiFineInnerPos, roiFineInner, roiFineOuter, roiCourse, roiCenter = roi
	
	cropFineOuter = crop(frame, roiFineOuter)
	cropCoarse = crop(frame, roiCourse)
	#gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
	
	dilation = cv2.dilate(cropFineOuter, dKernel, iterations=1)
	ret,thresh = cv2.threshold(dilation,100,255,cv2.THRESH_BINARY)

	replaceCrop(frame, roiFineOuter, thresh)
	
	meanCourse = pow(cropCoarse.mean(), 1)
	mean = 0
	direction = 1
	if(meanCourse > 10):
		# this could potentially be made more efficient by cropping from cropFineOuter
		cropFineInner = crop(frame, roiFineInner)
		# this could potentially be made more efficient by cropping from cropFineInner
		cropFineInnerNeg = crop(frame, roiFineInnerNeg)
		cropFineInnerPos = crop(frame, roiFineInnerPos)

		meanFine = pow(cropFineInner.mean(), exp)
		direction = np.sign(cropFineInnerPos.mean() - cropFineInnerNeg.mean())
	
		mean = meanFine
	
	distance = direction * (pow(255, exp) - mean)
	
	return distance


def drawRect(frame, points):
	cv2.rectangle(frame, points[0], points[1], (0, 255, 0))


def drawRoi(frame, roi):
	for rect in roi[2:5]:
		drawRect(frame, rect)
	center = roi[5]
	cv2.line(frame, (center[0] - 5, center[1]), (center[0] + 5, center[1]), (0, 255, 0), 1)
	cv2.line(frame, (center[0], center[1] - 5), (center[0], center[1] + 5), (0, 255, 0), 1)

def picameraToCVTrack():
	global roiX, roiY, moving, l1, l2, l3, w, selectedAxis, dilationKernel, calibrate, oscClient
	
	while True:
		frame = picam2.capture_buffer("lores")
		frame = frame[:s1 * h1].reshape((h1, s1))
		#frame = picam2.capture_array("lores")
		#frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)

		distanceX = track(frame, roiX, dilationKernel)
		distanceY = track(frame, roiY, dilationKernel)

		oscClient.send_message("/distanceX", distanceX)

		drawRoi(frame, roiX)
		drawRoi(frame, roiY)

		cv2.putText(frame, "{}: {:.2f}".format("distance x", distanceX), (10, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 3)
		cv2.putText(frame, "{}: {:.2f}".format("distance y", distanceY), (10, 200), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 3)

		if calibrate:
			cv2.imshow("Frame", frame)
			#cv2.imshow("Process", tresh)

		# Press Q on keyboard to exit
		key = cv2.waitKey(20)
		#if key == 32:
		#	cv2.waitKey()
		if key == ord('+'):
			dilationVal = dilationVal + 1
			dilationKernel = genDKernel(dilationVal)
		elif key == ord('-'):
			if dilationVal > 0:
				dilationVal = dilationVal - 1
				dilationKernel = genDKernel(dilationVal)
		elif key == ord('x'):
			selectedAxis = 'x'
		elif key == ord('y'):
			selectedAxis = 'y'
		elif key == ord('c'):
			if calibrate:
				calibrate = False
				cv2.destroyAllWindows()
			else:
				calibrate = True
				cv2.startWindowThread()
		#elif key == ord('q'):
		#	break


class TrackerThread(QThread):
	def __init__(self, target=None):
		super().__init__()
		self.target = target
		
	def run(self):
		if self.target:
			self.target()


class MainWindow(QGlPicamera2):
	def __init__(self, *args, **kwargs):
		super(MainWindow, self).__init__(*args, **kwargs)
		trackerThread = TrackerThread(target=picameraToCVTrack)
		trackerThread.start()
		self.shortcut_close_window = QShortcut(QKeySequence('f'), self)
		self.shortcut_close_window.activated.connect(self.goFullscreen)
		self.setWindowFlags(Qt.FramelessWindowHint)
		self.move(0, 0)

	def mousePressEvent(self, event):
		self.oldPos = event.globalPos()

	def mouseMoveEvent(self, event):
		delta = QPoint (event.globalPos() - self.oldPos)
		self.move(self.x() + delta.x(), self.y() + delta.y())
		self.oldPos = event.globalPos()

	def goFullscreen(self):
		if self.isFullScreen():
			#self.setWindowFlags(self._flags)
			self.showNormal()
		else:
			#self._flags = self.windowFlags()
			#self.setWindowFlags(Qt.WindowCloseButtonHint | Qt.WindowType_Mask)
			self.showFullScreen()

picam2 = Picamera2()
#picam2.start_preview(Preview.QTGL)
#max resolution is (4056, 3040) which is more like 10 fps
config = picam2.create_preview_configuration(main={"size": (2028, 1520)}, lores={"size": (1920, 1440), "format": "YUV420"})
picam2.configure(config)

app = QApplication([])
qpicamera2 = MainWindow(picam2, width=1920, height=1440, keep_ar=False)
qpicamera2.setWindowTitle("Qt Picamera2 App")

selectedAxis = 'x'
moving = False
l1 = 100
l2 = 300
l3 = 40
w = 20
roiXCenter = (960, 195)
roiYCenter = (615, 530)
roiX = rectsFromPoint(roiXCenter, l1, l2, l3, w, 'x')
roiY = rectsFromPoint(roiYCenter, l1, l2, l3, w, 'y')
dilationVal = 75
dilationKernel = genDKernel(dilationVal)
calibrate = True

oscClient = udp_client.SimpleUDPClient("127.0.0.1", 57120)

cv2.startWindowThread()
cv2.namedWindow("Frame")
cv2.setMouseCallback("Frame", moveROI)

(w0, h0) = picam2.stream_configuration("main")["size"]
(w1, h1) = picam2.stream_configuration("lores")["size"]
s1 = picam2.stream_configuration("lores")["stride"]

picam2.start()
qpicamera2.show()
app.exec()
