# startup:
# roscore -> start ros
# rosparam set cv_camera/device_id 0 -> set appropriate camera device
# rosrun cv_camera cv_camera_node -> start the camera
# roslaunch aruco_detect aruco_detect.launch camera:=cv_camera image:=image_raw dictionary:=16 transport:= fiducial_len:=0.1 # aruco marker detection
# python fiducial_to_2d_pos_angle.py -> compute position and angle of markers in 2d plane

import sys
import rospy
import pygame
import numpy as np
import socket
import scipy.integrate
import copy

import threading
from copy import deepcopy

import matplotlib.pyplot as plt
import matplotlib.animation as anim
import matplotlib.patches as patch

from shapely.geometry import Polygon

import time

from casadi_opt import OpenLoopSolver

from marker_pos_angle.msg import id_pos_angle

from collections import OrderedDict

from argparse import ArgumentParser

MSGLEN = 32
def myreceive(sock):
    chunks = []
    bytes_recd = 0
    while bytes_recd < MSGLEN:
        chunk = sock.recv(min(MSGLEN - bytes_recd, 2048))
        if chunk == b'':
            raise RuntimeError("socket connection broken")
        chunks.append(chunk)
        bytes_recd = bytes_recd + len(chunk)

        if chunk[-1] == '\n':
            break

    return b''.join(chunks)


class Robot:
    def __init__(self, id, ip):
        self.pos = None
        self.orient = None
        self.grid_pos = (0,0,0)

        self.id = id

        self.pos = None
        self.euler = None

        self.ip = ip
        self.socket = socket.socket()

        # variables for measurements
        self.tms_0 = None
        self.xm_0 = None
        self.tms = None
        self.xms = None

        # currently active control
        self.u1 = 0.0
        self.u2 = 0.0


    def connect(self):
        # connect to robot
        try:
                self.socket.connect((self.ip, 1234))  # connect to robot
        except socket.error:
            print("could not connect to socket")
            sys.exit(1)

def f_ode(t, x, u):
    # dynamical model of the two-wheeled robot
    # TODO: find exact values for these parameters
    r = 0.03
    R = 0.05
    d = 0.02

    theta = x[2]

    omega_r = u[0]
    omega_l = u[1]

    dx = np.zeros(3)

    dx[0] = (r/2.0 * np.cos(theta) - r*d/(2.0*R) * np.sin(theta)) * omega_r \
          + (r/2.0 * np.cos(theta) + r*d/(2.0 * R) * np.sin(theta)) * omega_l
    dx[1] = (r/2.0 * np.sin(theta) + r*d/(2.0*R) * np.cos(theta)) * omega_r \
          + (r/2 * np.sin(theta) - r*d/(2.0*R) * np.cos(theta)) * omega_l
    dx[2] = -r/(2.0*R) * (omega_r - omega_l)

    return dx

class RemoteController:
    def __init__(self, id, ip):

        self.anim_stopped = False

        self.robots = [Robot(11, '192.168.1.101'), Robot(14, '192.168.1.103')]
        #self.robots = [Robot(15, '192.168.1.102')]
        #self.robots = [Robot(id, ip)]

        self.robot_ids = {}
        for r in self.robots:
            self.robot_ids[r.id] = r

        self.valid_cmds = ['forward', 'backward', 'turn left', 'turn right']

        # socket for movement commands
        self.comm_socket = socket.socket()
        self.comm_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        for robot in self.robots:
            robot.connect()

        #self.comm_socket.bind((socket.gethostname(), 1337))
        self.comm_socket.bind(('', 1337))
        self.comm_socket.listen(5)

        self.t = time.time()

        # variables for simulated state
        self.x0 = None
        self.ts = np.array([])
        self.xs = []


        # variable for mpc open loop
        self.ol_x = None
        self.ol_y = None

        self.mutex = threading.Lock()

        # ROS subscriber for detected markers
        marker_sub = rospy.Subscriber("/marker_id_pos_angle", id_pos_angle, self.measurement_callback)

        # pid parameters
        self.controlling = False

        # animation
        self.fig = plt.figure()
        self.ax = self.fig.add_subplot(2,2,1)
        self.ax2 = self.fig.add_subplot(2, 2, 2)
        self.ax3 = self.fig.add_subplot(2, 2, 4)

        self.xdata, self.ydata = [], []
        self.line, = self.ax.plot([],[], color='grey', linestyle=':')
        self.line_sim, = self.ax.plot([], [])
        self.line_ol, = self.ax.plot([],[], color='green', linestyle='--')
        self.dirm, = self.ax.plot([], [])
        self.dirs, = self.ax.plot([], [])

        self.line_x, = self.ax2.plot([],[])
        self.line_y, = self.ax3.plot([], [])

        self.track_line_inner, = self.ax.plot([], [])
        self.track_line_outer, = self.ax.plot([], [])

        self.ax.set_xlabel('x-position')
        self.ax.set_ylabel('y-position')
        self.ax.grid()

        self.ax2.set_xlabel('Zeit t')
        self.ax2.set_ylabel('x-position')
        self.ax2.grid()

        self.ax3.set_xlabel('Zeit t')
        self.ax3.set_ylabel('y-position')
        self.ax3.grid()

        self.mstep = 2
        self.ols = OpenLoopSolver(N=20, T=1.0)
        self.ols.setup()
        self.dt = self.ols.T / self.ols.N

        self.target = (0.0, 0.0, 0.0)

        # integrator
        self.r = scipy.integrate.ode(f_ode)
        self.omega_max = 5.0
        self.control_scaling = 0.3
        #self.omega_max = 13.32


    def ani(self):

        print("starting animation")
        self.ani = anim.FuncAnimation(self.fig, init_func=self.ani_init, func=self.ani_update, interval=10, blit=True)
        plt.ion()
        plt.show(block=True)

    def ani_init(self):
        self.ax.set_xlim(-2, 2)
        self.ax.set_ylim(-2, 2)
        self.ax.set_aspect('equal', adjustable='box')

        self.ax2.set_ylim(-2, 2)
        self.ax2.set_xlim(0, 10)

        self.ax3.set_ylim(-2, 2)
        self.ax3.set_xlim(0, 10)

        self.track_line_inner.set_data(self.track.inner_poly.exterior.xy)
        self.track_line_outer.set_data(self.track.outer_poly.exterior.xy)

        return self.line, self.line_sim, self.dirm, self.dirs, self.line_ol,\
               self.track_line_inner, self.track_line_outer, self.line_x,self.line_y,

    def ani_update(self, frame):
        if self.anim_stopped:
            self.ani.event_source.stop()
            sys.exit(0)

        #print("plotting")
        self.mutex.acquire()
        try:
            # copy data for plot from global arrays
            if self.tms is not None:
                tm_local = deepcopy(self.tms)
                xm_local = deepcopy(self.xms)

                if len(tm_local) > 0:
                    # plot path of the robot
                    self.line.set_data(xm_local[:,0], xm_local[:,1])

                    # compute and plot direction the robot is facing
                    a = xm_local[-1, 0]
                    b = xm_local[-1, 1]

                    a2 = a + np.cos(xm_local[-1, 2]) * 0.2
                    b2 = b + np.sin(xm_local[-1, 2]) * 0.2

                    self.dirm.set_data(np.array([a, a2]), np.array([b, b2]))

                    n_plot = 300
                    if len(tm_local) > n_plot:
                        # plot x and y coordinate
                        self.line_x.set_data(tm_local[-n_plot:] - (tm_local[-1] - 10), xm_local[-n_plot:,0])
                        self.line_y.set_data(tm_local[-n_plot:] - (tm_local[-1] - 10), xm_local[-n_plot:, 1])

            ts_local = deepcopy(self.ts)
            xs_local = deepcopy(self.xs)

            if len(ts_local) > 0:
                # plot simulated path of the robot
                self.line_sim.set_data(xs_local[:,0], xs_local[:,1])

                # compute and plot direction the robot is facing
                a = xs_local[-1, 0]
                b = xs_local[-1, 1]

                a2 = a + np.cos(xs_local[-1, 2]) * 0.2
                b2 = b + np.sin(xs_local[-1, 2]) * 0.2

                self.dirs.set_data(np.array([a, a2]), np.array([b, b2]))

            ol_x_local = deepcopy(self.ol_x)
            ol_y_local = deepcopy(self.ol_y)

            if ol_x_local is not None:
                self.line_ol.set_data(ol_x_local, ol_y_local)
            else:
                self.line_ol.set_data([],[])

        finally:
            self.mutex.release()

        return self.line, self.line_sim, self.dirm, self.dirs, self.line_ol, self.track_line_inner, self.track_line_outer,\
               self.line_x, self.line_y, 

    def measurement_callback(self, data):
        #print(data)

        # detect robots
        if data.id in self.robot_ids:
            r = self.robot_ids[data.id]

            r.pos = (data.x, data.y) # only x and y component are important for us
            r.euler = data.angle

            # save measured position and angle for plotting
            measurement = np.array([r.pos[0], r.pos[1], r.euler])
            if r.tms_0 is None:
                r.tms_0 = time.time()
                r.xm_0 = measurement

                self.mutex.acquire()
                try:
                    r.tms = np.array([0.0])
                    r.xms = measurement
                finally:
                    self.mutex.release()
            else:
                self.mutex.acquire()
                try:
                    r.tms = np.vstack((r.tms, time.time() - r.tms_0))
                    r.xms = np.vstack((r.xms, measurement))
                finally:
                    self.mutex.release()

    def controller(self):
        print("starting control")

        running = True
        while running:
            (clientsocket, address) = self.comm_socket.accept()

            connected = True
            while connected:
                try:
                    data = myreceive(clientsocket)
                    print(data)

                    try:
                        robot_id, cmd = data.split(',')
                        robot_id = int(robot_id)
                        cmd = cmd.strip()

                        if robot_id in self.robot_ids and cmd in self.valid_cmds:
                            self.mpc_control(robot_id, cmd)
                        elif cmd == 'quit':
                            clientsocket.close()
                            self.comm_socket.close()
                            connected = False
                            running = False
                        else:
                            print("invalid command or robot id!")
                    except ValueError:
                        print("could not process command!")

                except RuntimeError:
                    print("disconnected")
                    connected = False
                    clientsocket.close()

    def mpc_control(self, robot_id, cmd):
            robot = self.robot_ids[robot_id] # get robot to be controlled
            grid_pos = robot.grid_pos   # grid position of the robot

            # compute new grid position and orientation
            if cmd == 'forward':
                new_x = grid_pos[0] + 1 * np.cos(grid_pos[2])
                new_y = grid_pos[1] + 1 * np.sin(grid_pos[2])
                new_angle = grid_pos[2]
            elif cmd == 'backward':
                new_x = grid_pos[0] - 1 * np.cos(grid_pos[2])
                new_y = grid_pos[1] - 1 * np.sin(grid_pos[2])
                new_angle = grid_pos[2]
            elif cmd == 'turn left':
                new_x = grid_pos[0]
                new_y = grid_pos[1]
                new_angle = np.unwrap([0, grid_pos[2] + np.pi / 2])[1]
            elif cmd == 'turn right':
                new_x = grid_pos[0]
                new_y = grid_pos[1]
                new_angle = np.unwrap([0, grid_pos[2] - np.pi / 2])[1]
            else:
                print("unknown command!")
                sys.exit(1)

            grid_pos = (new_x, new_y, new_angle)
            print("new grid pos for robot {}: {}".format(robot_id, grid_pos))

            self.target = np.array((0.25 * grid_pos[0], 0.25 * grid_pos[1], grid_pos[2]))

            self.pid = False
            self.mpc = True

            near_target = 0

            while near_target < 5:
                # open loop controller
                events = pygame.event.get()

                for event in events:
                    if event.type == pygame.KEYDOWN:
                        if event.key == pygame.K_UP:
                            self.controlling = True
                            self.t = time.time()
                        elif event.key == pygame.K_DOWN:
                            self.controlling = False
                            if self.robot_ids[robot_id].socket:
                                self.robot_ids[robot_id].socket.send('(0.0,0.0)\n')
                        elif event.key == pygame.K_0:
                            self.target = np.array([0,0,0])
                        elif event.key == pygame.K_PLUS:
                            self.control_scaling += 0.1
                            self.control_scaling = min(self.control_scaling, 1.0)
                            print("control scaling = ", self.control_scaling)
                        elif event.key == pygame.K_MINUS:
                            self.control_scaling -= 0.1
                            self.control_scaling = max(self.control_scaling, 0.1)
                            print("control scaling = ", self.control_scaling)
                        elif event.key == pygame.K_ESCAPE:
                            print("quit!")
                            self.controlling = False
                            if self.robot_ids[robot_id].socket:
                                self.robot_ids[robot_id].socket.send('(0.0,0.0)\n')
                            self.anim_stopped = True
                            return
                        elif event.key == pygame.QUIT:
                            print("quit!")
                            self.controlling = False
                            if self.robot_ids[robot_id].socket:
                                self.robot_ids[robot_id].socket.send('(0.0,0.0)\n')
                            self.anim_stopped = True
                            return

                if self.mpc:
                    x_pred = self.get_measurement_prediction(robot_id)

                    tmpc_start = time.time()

                    error_pos = np.linalg.norm(x_pred[0:2] - self.target[0:2])
                    angles_unwrapped = np.unwrap([x_pred[2], self.target[2]])   # unwrap angle to avoid jump in data
                    error_ang = np.abs(angles_unwrapped[0] - angles_unwrapped[1])
                    #print("error pos = ", error_pos)
                    print("error_pos = {}, error_ang = {}".format(error_pos, error_ang))

                    if error_pos > 0.075 or error_ang > 0.35:
                        # solve mpc open loop problem
                        res = self.ols.solve(x_pred, self.target)

                        #us1 = res[0]
                        #us2 = res[1]
                        us1 = res[0] * self.control_scaling
                        us2 = res[1] * self.control_scaling
                        #print("u = {}", (us1, us2))

                        # save open loop trajectories for plotting
                        self.mutex.acquire()
                        try:
                            self.ol_x = res[2]
                            self.ol_y = res[3]
                        finally:
                            self.mutex.release()

                        tmpc_end = time.time()
                        #print("---------------- mpc solution took {} seconds".format(tmpc_end - tmpc_start))
                        dt_mpc = time.time() - self.t
                        if dt_mpc < self.dt: # wait until next control can be applied
                            #print("sleeping for {} seconds...".format(self.dt - dt_mpc))
                            time.sleep(self.dt - dt_mpc)
                    else:
                        us1 = [0] * self.mstep
                        us2 = [0] * self.mstep

                        near_target += 1
                        robot.grid_pos = grid_pos

                    # send controls to the robot
                    for i in range(0, self.mstep): # option to use multistep mpc if len(range) > 1
                        u1 = us1[i]
                        u2 = us2[i]
                        if self.robot_ids[robot_id].socket:
                            self.robot_ids[robot_id].socket.send('({},{})\n'.format(u1, u2))
                            if i < self.mstep:
                                time.sleep(self.dt)
                        self.t = time.time() # save time the most recent control was applied


    def get_measurement_prediction(self, robot_id):
        # get measurement
        self.mutex.acquire()
        try:
            window = 3
            last_measurement = copy.deepcopy(self.robot_ids[robot_id].xms[-window:])
            #print("last_measurements = {}".format(last_measurement))
            #print("mean = {}".format(np.mean(last_measurement, axis=0)))
            last_measurement = np.mean(last_measurement, axis=0)
            last_time = copy.deepcopy(self.robot_ids[robot_id].tms[-1])
        finally:
            self.mutex.release()

        # prediction of state at time the mpc will terminate
        self.r.set_f_params(np.array([self.robot_ids[robot_id].u1 * self.omega_max, self.robot_ids[robot_id].u2 * self.omega_max]))

        self.r.set_initial_value(last_measurement, last_time)

        x_pred = self.r.integrate(self.r.t + self.dt)

        return x_pred

    def get_measurement(self):
        self.mutex.acquire()
        try:
            last_measurement = copy.deepcopy(self.xms[-1:])
        finally:
            self.mutex.release()
        return last_measurement[0]

    def pos_getter(self):
        while True:
            x_pred = self.get_measurement_prediction()

            print("pos = ", x_pred)

def main(args):
    parser = ArgumentParser()
    parser.add_argument('id', metavar='id', type=str, help='marker id of the controlled robot')
    parser.add_argument('ip', metavar='ip', type=str, help='ip address of the controlled robot')
    args = parser.parse_args()

    marker_id = int(args.id)
    ip = args.ip


    rospy.init_node('controller_node', anonymous=True)

    rc = RemoteController(marker_id, ip)

    pygame.init()

    screenheight = 480
    screenwidth = 640
    pygame.display.set_mode([screenwidth, screenheight])

 #   print("waiting until track is completely detected")
 #   while not rc.track.track_complete:
 #       pass

    #threading.Thread(target=rc.input_handling).start()
    controller_thread = threading.Thread(target=rc.controller)
    controller_thread.start()

    #time.sleep(10)
    #rc.ani()


if __name__ == '__main__':
    main(sys.argv)