#!/usr/bin/env python3
'''
Meant to run on your PC.
Receives data generated by rmServerWacomInput.py,
moves the mouse and presses accordingly.
Configure area below!
'''


from sys import argv
import socket
import struct
from pynput.mouse import Button, Controller

from kalman_filter import KalmanFilter
import numpy as np

mouse = Controller()

def mouseMoveAbs(x, y):
	'''The 'mouse.move()' method only moves relative.
	   This funtion works with absolute values.'''
	pos = mouse.position
	mouse.move(x - pos[0], y - pos[1])

# ----------
# Config:

ONLY_DEBUG=False  # Only show data. Don't move mouse

# Area on your display (remember to keep correct ratio (1872:1404 or 312:234) or your input will get streched/squashed!)
SCREEN_DRAW_AREA_FROM_X = 0
SCREEN_DRAW_AREA_FROM_Y = 0
SCREEN_DRAW_AREA_TO_X = 1700  # Ratio will match roughly but not exact!
SCREEN_DRAW_AREA_TO_Y = 1000  # ↑
SCREEN_DRAW_BUTTON_PRESSURE = 1  # Pressure needed to click the left mouse button (0 contact; 4095 = hardest)

# ----------

WACOM_WIDTH = 15725   # Values just checked by drawing to the edges
WACOM_HEIGHT = 20967  # ↑

screenRectWidth = SCREEN_DRAW_AREA_TO_X - SCREEN_DRAW_AREA_FROM_X
screenRectHeight = SCREEN_DRAW_AREA_TO_Y - SCREEN_DRAW_AREA_FROM_Y
ratioX = screenRectWidth / WACOM_WIDTH
#ratioY = screenRectHeight / WACOM_HEIGHT
ratioY = ratioX

# Source: https://github.com/canselcik/libremarkable/blob/master/src/input/wacom.rs
EV_SYNC = 0
EV_KEY = 1
EV_ABS = 3
WACOM_EVCODE_PRESSURE = 24
WACOM_EVCODE_DISTANCE = 25
WACOM_EVCODE_XTILT = 26
WACOM_EVCODE_YTILT = 27
WACOM_EVCODE_XPOS = 0
WACOM_EVCODE_YPOS = 1

lastXPos = -1
lastYPos = -1
lastXTilt = -1
lastYTilt = -1
lastDistance = -1
lastPressure = -1

mouseButtonPressed = False

client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
client.connect(('10.11.99.1' if len(argv) == 1 else argv[1], 33333))

# Source: https://github.com/DAA233/kalman-filter/
stateMatrix = np.zeros((4, 1), np.float32)  # [x, y, delta_x, delta_y]
estimateCovariance = np.eye(stateMatrix.shape[0])
transitionMatrix = np.array([[1, 0, 1, 0],[0, 1, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32)
processNoiseCov = np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]], np.float32) * 0.001
measurementStateMatrix = np.zeros((2, 1), np.float32)
observationMatrix = np.array([[1,0,0,0],[0,1,0,0]],np.float32)
measurementNoiseCov = np.array([[1,0],[0,1]], np.float32) * 1
kalman = KalmanFilter(X=stateMatrix,
                      P=estimateCovariance,
                      F=transitionMatrix,
                      Q=processNoiseCov,
                      Z=measurementStateMatrix,
                      H=observationMatrix,
                      R=measurementNoiseCov)

while True:
	evDevType, evDevCode, evDevValue = struct.unpack('HHi', client.recv(8))
	if evDevType == EV_ABS:
		if evDevCode == WACOM_EVCODE_XPOS:
			lastYPos = evDevValue  # X is Y
		elif evDevCode == WACOM_EVCODE_YPOS:
			lastXPos = evDevValue  # Y is X
		elif evDevCode == WACOM_EVCODE_XTILT:
			lastXTilt = evDevValue
		elif evDevCode == WACOM_EVCODE_YTILT:
			lastYTilt = evDevValue
		elif evDevCode == WACOM_EVCODE_DISTANCE:
			lastDistance = evDevValue
		elif evDevCode == WACOM_EVCODE_PRESSURE:
			if not ONLY_DEBUG:
				if not mouseButtonPressed and evDevValue > SCREEN_DRAW_BUTTON_PRESSURE:
					mouse.press(Button.left)
					mouseButtonPressed = True
				elif mouseButtonPressed and evDevValue <= SCREEN_DRAW_BUTTON_PRESSURE:
					mouse.release(Button.left)
					mouseButtonPressed = False

			lastPressure = evDevValue

		if ONLY_DEBUG:
			print('XPos: %5d | YPos: %5d | XTilt: %5d | YTilt: %5d | Distance: %3d | Pressure: %4d' % (lastXPos, lastYPos, lastXTilt, lastYTilt, lastDistance, lastPressure))
		else:
			screenY = SCREEN_DRAW_AREA_FROM_X + (WACOM_WIDTH - lastXPos) * ratioX                 	# (X doesn't need to invert)
			screenX = SCREEN_DRAW_AREA_FROM_Y + (WACOM_HEIGHT - lastYPos) * ratioY  # (Y has to be inverted)

			currentMeasurement = np.array([[np.float32(screenX)], [np.float32(screenY)]])
			currentPrediction = kalman.predict()

			cmx, cmy = currentMeasurement[0], currentMeasurement[1]
			cpx, cpy = currentPrediction[0], currentPrediction[1]

			kalman.correct(currentMeasurement)

			mouseMoveAbs(int(np.round(cpx)), int(np.round(cpy)))

