import cv2, sys
import numpy as np
import tkinter as tk
from tkinter import ttk
from PIL import ImageTk, Image

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

from vision import process, obtain_depth_view

try:
	import rclpy
	from rclpy.node import Node
	import std_msgs.msg
	import sensor_msgs.msg
	import lfr_msgs.msg
	import threading

	_ROS2_FOUND = True
except ImportError:
	print("Could not import ROS2; please activate!")
	_ROS2_FOUND = False

class Dashboard:
	def __init__(self):
		self.root = tk.Tk()
		self.root.geometry("900x600")

		self.viewpanel = tk.Frame(self.root)
		self.right = tk.Frame(self.root)

		self.DEPTH_SIZE = 300
		self.PLOT_SIZE = 600
		self.REDSENSOR_HISTORY = 100
		self.INTEGRATION_INTERVAL = 400  # ms

		self.__finished = False

		self.state = {
			"redsensor": [],
			"handbrake": True,
			"sys_history": [],
			"sim_history": [],
		}

		# Show the depth image
		self.image_notebook = ttk.Notebook(self.right)
		self.depth_image_canvas = tk.Canvas(self.image_notebook, width=self.DEPTH_SIZE, height=self.DEPTH_SIZE)
		self.depth_image_canvas.pack(side="top")
		self.depth_image = None
		self.dit = self.depth_image_canvas.create_text(self.DEPTH_SIZE // 2, self.DEPTH_SIZE // 2,
		                                               text="Waiting for Signal...", anchor=tk.CENTER)
		self.di = self.depth_image_canvas.create_image(0, 0, anchor=tk.NW)

		self.processed_image_canvas = tk.Canvas(self.image_notebook, width=self.DEPTH_SIZE, height=self.DEPTH_SIZE)
		self.processed_image_canvas.pack(side="top")
		self.processed_image = None
		self.pit = self.processed_image_canvas.create_text(self.DEPTH_SIZE // 2, self.DEPTH_SIZE // 2,
		                                                   text="Waiting for Signal...", anchor=tk.CENTER)
		self.pi = self.processed_image_canvas.create_image(0, 0, anchor=tk.NW)

		self.image_notebook.add(self.depth_image_canvas, text="Depth Vision Image")
		self.image_notebook.add(self.processed_image_canvas, text="OpenCV Processed Image")
		self.image_notebook.pack(side="top", padx=5, pady=2)

		# Show the color sensor value (+ history plot)
		self.colvalue_label = tk.Label(self.right, text="Sensor Value: ???")
		self.colvalue_label.pack(side="top", padx=5, pady=2)
		self.colfig = plt.figure(figsize=(3, 1), dpi=100)
		self.colfig_ax = self.colfig.add_subplot(111)
		self.colfig_ax.set_xlim((0, self.REDSENSOR_HISTORY))
		self.colfig_ax.set_ylim((0, 1))
		self.colfig_line, = self.colfig_ax.plot([], [], c="red")
		self.colfig_canvas = FigureCanvasTkAgg(self.colfig, master=self.right)
		self.colfig_canvas.draw()
		self.colfig_canvas.get_tk_widget().pack(side="top", padx=5, pady=2)
		self.brake_button = tk.Button(self.right, text="Toggle Handbrake Off", command=self.toggle_brake)
		self.brake_button.pack(side="top", fill="x", padx=5, pady=2)

		# Create the drawing canvas
		self.canvas = tk.Canvas(self.viewpanel, width=self.PLOT_SIZE, height=self.PLOT_SIZE)
		self.fig = plt.figure(figsize=(5, 5), dpi=100)
		self.fig_ax = self.fig.add_subplot(111)
		self.fig_ax.set_xlim((0.2, 0.9))
		self.fig_ax.set_ylim((0.1, 0.8))
		self.fig_sys_line, = self.fig_ax.plot([], [], c="red", marker='.', alpha=.5)
		self.fig_sim_line, = self.fig_ax.plot([], [], c="blue", marker='.', alpha=.5)
		self.canvas = FigureCanvasTkAgg(self.fig, master=self.viewpanel)
		self.canvas.draw()
		self.canvas.get_tk_widget().pack(side="top", fill="both", expand="yes")

		self.viewpanel.pack(side="left", fill="both", expand="yes")
		self.right.pack(side="right", fill="both", expand="no")
		self.root.wm_title("AGV Detection")
		self.root.wm_protocol("WM_DELETE_WINDOW", self.on_close)

		# ROS2
		if _ROS2_FOUND:
			self.node = Node("dashboardLFR")
			self.state_publisher = self.node.create_publisher(lfr_msgs.msg.State, "lfr/state_sys", 10)
			self.brake_publisher = self.node.create_publisher(std_msgs.msg.Bool, "lfr/handbrake", 10)
			self.create_listeners()
			self.spinner = threading.Thread(target=lambda: rclpy.spin(self.node), daemon=True)
			self.spinner.start()

	def done(self):
		return self.__finished

	def on_close(self):
		self.__finished = True
		if _ROS2_FOUND:
			self.node.destroy_node()
			self.spinner.join(1)
			self.root.quit()
		else:
			self.root.quit()

	def create_listeners(self):
		self.node.create_subscription(sensor_msgs.msg.Image, "lfr/depth_processor", self.update_depthimage, 10)
		self.node.create_subscription(sensor_msgs.msg.Image, "lfr/depth_processor", self.update_processimage, 10)
		self.node.create_subscription(std_msgs.msg.Float32, "lfr/redsensor", self.update_redsensor, 10)
		self.node.create_subscription(lfr_msgs.msg.State, "lfr/state_sys", self.update_system, 10)
		self.node.create_subscription(std_msgs.msg.Bool, "lfr/handbrake", self.update_handbrake, 10)
		self.colfig_ani = animation.FuncAnimation(self.colfig, lambda _: self.update_redsensor_plot(), interval=self.INTEGRATION_INTERVAL)
		self.fig_ani = animation.FuncAnimation(self.fig, lambda _: self.update_plot(), interval=self.INTEGRATION_INTERVAL)

	def update_redsensor(self, msg):
		self.state["redsensor"].append(msg.data)
		self.colvalue_label.config(text="Sensor Value: %.5f" % msg.data)

	def update_system(self, msg):
		state = msg.x, msg.y, msg.w
		self.state["sys_history"].append(state)

	def update_handbrake(self, msg):
		self.state["handbrake"] = msg.data
		self.brake_button.config(text="Toggle Handbrake %s" % ("Off" if msg.data else "On"))

	def toggle_brake(self):
		if _ROS2_FOUND and rclpy.ok():
			msg = std_msgs.msg.Bool()
			msg.data = not self.state["handbrake"]
			self.brake_publisher.publish(msg)

	def update_redsensor_plot(self):
		while len(self.state["redsensor"]) > self.REDSENSOR_HISTORY:
			self.state["redsensor"].pop(0)
		data = self.state["redsensor"][:]
		self.colfig_line.set_data(list(range(len(data))), data)

	def update_plot(self):
		system = self.state["sys_history"][:]
		xs = [x for x, _, __ in system]
		ys = [1. - y for _, y, __ in system]
		self.fig_sys_line.set_data(xs, ys)

	def update_depthimage(self, msg):
		if self.dit is not None:
			self.depth_image_canvas.delete(self.dit)
			self.dit = None
		dt = np.dtype('float32')
		dt = dt.newbyteorder('<')
		current_frame = np.ndarray(shape=(msg.height, msg.width, 1), dtype=dt, buffer=bytearray(msg.data))

		im_color = obtain_depth_view(current_frame, self.DEPTH_SIZE)

		self.depth_image = Image.fromarray(im_color)
		self.depth_image = ImageTk.PhotoImage(self.depth_image)

		self.depth_image_canvas.itemconfig(self.di, image=self.depth_image)

	def update_processimage(self, msg):
		if self.pit is not None:
			self.processed_image_canvas.delete(self.pit)
			self.pit = None
		dt = np.dtype('float32')
		dt = dt.newbyteorder('<')
		current_frame = np.ndarray(shape=(msg.height, msg.width, 1), dtype=dt, buffer=bytearray(msg.data))

		new_frame, state = process(current_frame.copy(), self.DEPTH_SIZE)

		# show processed image
		self.processed_image = Image.fromarray(new_frame)
		self.processed_image = ImageTk.PhotoImage(self.processed_image)

		self.processed_image_canvas.itemconfig(self.pi, image=self.processed_image)

		# send identified state via ROS2
		if _ROS2_FOUND and rclpy.ok():
			q = lfr_msgs.msg.State()
			q.x = float(state[0])
			q.y = float(state[1])
			q.w = float(state[2])
			self.state_publisher.publish(q)


if __name__ == '__main__':
	rclpy.init()

	app = Dashboard()
	app.root.mainloop()

	# TODO: somehow this hangs the termination of the program..
	# rclpy.shutdown()