import serial
import time
import threading
import math


print("PrintBuddy V0.1")
print("Made by Thomas Van Acker")
print("")

# Static var init
COMMANDS = {
	"help":"List all commands with explanation",
	"exit":"Disconnect from printer and exit",
	"raw":"Send raw gcode commands",
	"home":"Auto-home the machine",
	"load":"Load a gcode file",
	"add":"Manipulate the loaded gcode",
	"print":"Print the loaded gcode",
	"level":"Start bed leveling wizard",
	"mesh":"Start mesh bed leveling wizard"
}

GCODE_FILCH_LAYER = ["G28 X Y", "M84", "M300 P3000", "M0", "M600", "M400", "G28 X Y", "G0 F1500"]

# Function declarations
def send(c, silent=False):
	global s
	if not silent:
		print("%s"%c)
	c += "\n"
	s.write(c.encode("utf-8"))
	while True:
		r = s.readline().decode("utf-8").strip()
		if not silent:
			print(" > %s"%r)
		if r[0:2] == "ok":
			break

def threadPrintInput():
	# Checks for input from user during printing process
	global abortPrint
	while printing:
		i = input()
		
		if i == "stop" or i == "exit":
			abortPrint = True
			break

def meshZ(X, Y):
	# Interpolates values of mesh
	x = X/300*(MESH_SIZE-1)
	y = Y/300*(MESH_SIZE-1)

	xmin = math.floor(x)
	xmax = math.ceil(x)
	ymin = math.floor(y)
	ymax = math.ceil(y)

	z0 = mesh[ymin][xmin] + (mesh[ymin][xmax]-mesh[ymin][xmin])*(x-xmin)
	z1 = mesh[ymax][xmin] + (mesh[ymax][xmax]-mesh[ymax][xmin])*(x-xmin)

	z = z0 + (z1-z0)*(y-ymin)
	return z
	
# Var init
serialPort = "/dev/ttyUSB0"
serialBaud = 115200

gcodeRaw = ""
gcodeLines = []
gcodeLayers = 1
printing = False
abortPrint = False

MESH_SIZE = 9		# Number of points of a side in the mesh.
MESH_OFFSET = 1.00	# The distance above Z=0 where printbed is located. Mesh needs to be added to this number to get the absolute Z height.
mesh = [[0 for _ in range(MESH_SIZE)] for __ in range(MESH_SIZE)]
meshPositions = [["X%d Y%d"%(280/(MESH_SIZE-1)*x+10, 280/(MESH_SIZE-1)*y+10) for x in range(MESH_SIZE)] for y in range(MESH_SIZE)]

levelingCornersPos = {0:"X10 Y45",1:"X10 Y260",2:"X280 Y260",3:"X280 Y45"}


# Read mesh file
print("Reading mesh file...")
try:
	f = open("mesh.txt", "r")
	t = f.read()[:-1]
	f.close()

	lines = t.split("\n")
	MESH_SIZE = int(lines[0])
	mesh = [[0 for _ in range(MESH_SIZE)] for __ in range(MESH_SIZE)]
	meshPositions = [["X%d Y%d"%(280/(MESH_SIZE-1)*x+10, 280/(MESH_SIZE-1)*y+10) for x in range(MESH_SIZE)] for y in range(MESH_SIZE)]

	for i in range(MESH_SIZE*MESH_SIZE):
		mesh[int(i/MESH_SIZE)][int(i%MESH_SIZE)] = float(lines[i+1])

	print("Mesh loaded.")
except Exception as e:
	print("Couldn't read mesh file. Using default mesh.")
	print(e)
print("Current mesh:")
for i in mesh:
	for j in i:
		print("{:5.2f}".format(j), end=" ")
	print()

# Connect to printer
print("\nConnecting to printer...")
s = serial.Serial(serialPort, serialBaud)
time.sleep(1)
s.write(b"\r\n\r\n") # Wake printer
time.sleep(2)
readBytes = s.inWaiting()
r = s.read(readBytes).decode("utf-8")
print(" > %s"%r)
time.sleep(5)
print("Connected!")

# Start main loop
while True:
	print("\n# Please enter a command ('help' for command list)")
	cmd = input()

	if cmd == "help":
		# List all commands
		print("\nCOMMAND\t\tACTION")
		for k in COMMANDS.keys():
			print("%s\t\t%s"%(k, COMMANDS[k]))

	elif cmd == "home":
		# Home machine
		send("G28", silent=True) #Auto-home
		send("M84", silent=True) #Disable steppers

	elif cmd == "raw":
		# Send raw gcode commands
		print("Entering raw mode. Proceed at your own risk. Type 'exit' to exit raw mode.")
		
		while True:
			print("\n# RAW: Please enter a gcode command (or 'exit' to exit)");
			r = input()
			if r == "exit":
				break
			send(r)

		print("Exiting raw mode.")

	elif cmd == "load":
		# Load gcode file
		print("# Please enter the path to the gcode file you want to load")
		path = input()
		
		# Try to open and parse file
		try:
			print("Trying to open file...")
			f = open(path, "r")
			print("Reading...")
			gcodeRaw = f.read()[:-1]
			f.close()
			print("Parsing...")
			gcodeLines = gcodeRaw.split("\n")

			# Iterate through all lines to get info
                        gcodeLines = 0
			for l in range(len(gcodeLines)):
				line = gcodeLines[l]

				if line[0] == ";":
					# When comment
					if line.startswith(";LAYER:"):
						l = int(line[line.index(":") +1:])
                                                if l > gcodeLayers:
                                                    gcodeLayers = l

			print("#Layers: %d"%gcodeLayers)
			print("%d lines"%len(gcodeLines))

			print("File loaded! Use 'add' to manipulate this gcode.")
		except Exception as e:
			print("Couldn't read file:")
			print(e)

			# Reset gcode vars
			gcodeRaw = ""
			gcodeLines = []
			gcodeLayers = 1

	elif cmd == "add":
		# Manipulate gcode

		# Check if gcode present
		if gcodeRaw == "" or len(gcodeLines) == 0:
			print("No gcode loaded. Use 'load' to load gcode.")
			continue

		# Display all commands
		print("COMMAND\t\tEXPLANATION")
		print("exit\t\tReturn")
		print("mesh\t\tEnable Mesh Bed Leveling")
		print("filch\t\tFilament change at specific layer")
		print("")
		print("# Please enter one of the above actions.")

		# Getting input
		act = input()
		
		if act == "exit":
			print("Exiting.")
			pass
		
		elif act == "mesh":
			# Add Mesh Bed Leveling
			print("Recalculating Z with Mesh Bed Leveling...")
			layerZ = 0
			layerIndex = 0
			for l in range(len(gcodeLines)):
				line = gcodeLines[l]

				# Check if comment
				if line[0] == ";":
					# Check if new layer
					if line.startswith(";LAYER:"):
						layerIndex = int(line[line.index(":") +1:])

				if not (line.startswith("G0") or line.startswith("G1")):
					continue

				# When this is a movement line
				tokens = line.split(" ")
				indexX = -1
				indexY = -1
				indexZ = -1
				for t in range(len(tokens)):
					if tokens[t][0] == "X":
						indexX = t
					elif tokens[t][0] == "Y":
						indexY = t
					elif tokens[t][0] == "Z":
						indexZ = t

				# Check if Z is given
				if not indexZ == -1:
					# When Z is set -> save as layer height
					layerZ = float(tokens[indexZ][1:])

				# Add Z if X or Y are set
				if indexX == -1 or indexY == -1:
					continue

				X = float(tokens[indexX][1:])
				Y = float(tokens[indexY][1:])
				Z = layerZ + MESH_OFFSET + (meshZ(X, Y)*max(0, float(10-layerIndex)/10.0) if layerIndex < 10 else 0)
				
				# Add Z coordinate
				tokens.append("Z{:.3f}".format(Z))
				newLine = " ".join(tokens)
				gcodeLines[l] = newLine
			print("Done!")

		elif act == "filch":
			# Filament change at layer
			layer = 0
			while True:
				# Ask for layer
				print("# FILCH: At which layer do you want to change filament?")
				l = input()

				# Check if valid
				try:
					layer = int(l)
					
					if layer > 0 and layer <= gcodeLayers:
						break

					print("Invalid.")
				except:
					print("Invalid.")

			if layer == 0:
				continue # exit

			# Search for start of this layer
			print("Searching for layer...")
			success = False
			for l in range(len(gcodeLines)):
				line = gcodeLines[l]

				if line[0] == ";":
					if line.startswith(";LAYER:"):
						thisLayer = int(line[line.find(":")+1:])

						# Check if right layer
						if thisLayer == layer:
							# Add filch gcode
							print("Adding gcode for filch...")
							gcodeLines[l:l] = GCODE_FILCH_LAYER
							success = True
							break
			print("Done." if success else "Couldn't find layer.")

		else:
			# When add action not found
			print("Command not found.")

	elif cmd == "print":
		# Print loaded gcode
		if len(gcodeLines) == 0:
			print("No file loaded yet. Use 'load' to load gcode.")
		else:
			# Start print
			printing = True
			abortPrint = False
			print("Print started. Type 'stop' at any time to abort the print.")
			send("M75", silent=True) # Start timer

			# Start thread to check for commands
			threadInput = threading.Thread(target=threadPrintInput, daemon=True)
			threadInput.start()

			# Send gcode line by line to printer
			for l in range(len(gcodeLines)):
				line = gcodeLines[l]

				# Check if needs to abort print
				if abortPrint:
					print("Aborting print!")
					break
				
				# Check if comment
				if line[0] == ";":
					# Check if new layer
					if line.startswith(";LAYER:"):
						layer = int(line[line.index(":") +1:])
						print("Printing layer %d/%d"%(layer, gcodeLayers))
				else:
					# Send line to printer
					send(line, silent=True)

			# Stop input thread
			printing = False

			# Stop print
			send("M104 S0", silent=True) # Hotend off
			send("M140 S0", silent=True) # Bed off
			send("M107", silent=True) # Fan off
			send("M77", silent=True) # Stop timer
			send("M31", silent=True) # Display time on lcd
			send("G28 X Y", silent=True) # Home X and Y
			send("G0 Y300 F1500", silent=True) # Move platform to user
			send("M84", silent=True) # Disable steppers
			send("M400", silent=True) # Wait for previous commands to finish
			send("M300 P3000", silent=True) # Beep
			print("Print done!")

	elif cmd == "level":
		# Normal leveling wizard
		print("Starting leveling wizard...")
		send("G28", silent=True) #Home
		send("G0 F7500", silent=True) # Set feedrate
		
		print("Please adjust the knob under the nozzle so that a piece of paper can slide under the nozzle with just a little bit of friction.")
		
		# Start corner loop
		i = 0
		while True:
			# Move to corner
			print("Moving to corner %d..."%i)
			send("G0 Z10", silent=True)
			send("G0 %s"%levelingCornersPos[i], silent=True)
			send("G0 Z0", silent=True)
			send("M84", silent=True) #Disable steppers

			# Prompt user
			print("Press ENTER to move to next corner. Type 'exit' to exit wizard.")
			r = input()

			if r == "exit":
				break

			# Next corner
			i = (i+1)%4

		# End wizard
		send("G0 F1500", silent=True) #Set slower feedrate
		send("G28", silent=True) #Home
		send("M84", silent=True) #Disable steppers


	elif cmd == "mesh":
		# Mesh leveling wizard
		print("Starting mesh leveling wizard...")
		send("G28", silent=True) #Home
		send("G0 F7500", silent=True) # Set feedrate
		
		print("\nSTEP 1: Normal leveling")
		print("Please adjust the knob under the nozzle so that a piece of paper can slide under the nozzle with just a little bit of friction.")
		
		# Start corner loop
		i = 0
		stop = False
		while True:
			# Move to corner
			print("Moving to corner %d..."%i)
			send("G0 Z10", silent=True)
			send("G0 %s"%levelingCornersPos[i], silent=True)
			send("G0 Z{:.2f}".format(MESH_OFFSET), silent=True)

			# Prompt user
			print("Press ENTER to move to next corner. Type 'next' to go to step 2. Type 'exit' to exit wizard.")
			r = input()

			if r == "exit":
				stop = True
				break
			elif r == "next":
				break

			# Next corner
			i = (i+1)%4

		# Step 2: get mesh deltas
		if not stop:
			print("\nSTEP 2: Mesh leveling")
			print("For every mesh point, lower the nozzle using the '+' and '-' buttons until the nozzle is at the right height. Then, enter 'next' to go to the next mesh point. Type 'exit' to exit the wizard.")

			# Move through mesh
			for y in range(MESH_SIZE):
				for x in range(MESH_SIZE):
					Z = 1.00	# Offset on top of MESH_OFFSET
					print("\nMoving to mesh point x=%d, y=%d"%(x,y))
					send("G0 Z10", silent=True)
					send("G0 %s"%meshPositions[y][x], silent=True)
					send("G0 Z{:.3f}".format(Z+MESH_OFFSET), silent=True)
					
					# Input loop
					stop = False
					while True:
						i = input()
						
						if i == "+":
							Z += 0.01
							send("G0 Z{:.3f}".format(Z+MESH_OFFSET), silent=True)
						elif i == "-":
							Z -= 0.01
							send("G0 Z{:.3f}".format(Z+MESH_OFFSET), silent=True)
						elif i == "++":
							Z += 0.1
							send("G0 Z{:.3f}".format(Z+MESH_OFFSET), silent=True)
						elif i == "--":
							Z -= 0.1
							send("G0 Z{:.3f}".format(Z+MESH_OFFSET), silent=True)
						elif i == "+++":
							Z += 0.5
							send("G0 Z{:.3f}".format(Z+MESH_OFFSET), silent=True)
						elif i == "---":
							Z -= 0.5
							send("G0 Z{:.3f}".format(Z+MESH_OFFSET), silent=True)
						elif i == "exit":
							stop = True
							break
						elif i == "next":
							mesh[y][x] = Z
							print("Mesh value saved (Z={:.3f})".format(Z))
							break
					
					# Check if needs to exit
					if stop:
						break
				# Check if needs to exit
				if stop:
					break

		# Print mesh
		print("\nCurrent Mesh:")
		for i in mesh:
			for j in i:
				print("{:6.3f}".format(j), end=" ")
			print()


		# Save mesh to file
		print("Writing mesh to file...")
		try:
			t = ""
			t += "%d\n"%MESH_SIZE
			for y in range(MESH_SIZE):
				for x in range(MESH_SIZE):
					t += "{:.3f}\n".format(mesh[y][x])
			f = open("mesh.txt", "w")
			f.write(t)
			f.close()
		except Exception as e:
			print("Couldn't write mesh to file.")
			print(e)
			

		# End wizard
		send("G0 F1500", silent=True) #Set slower feedrate
		send("G28", silent=True) #Home
		send("M84", silent=True) #Disable steppers


	elif cmd == "exit":
		break
	else:
		print("Command not found. Enter 'help' to list all commands.")
		


# Close serial
print("Disconnecting from printer...")
s.close()

# End of script
print("Thank you for using PrintBuddy!")