from collections import namedtuple
from fast_motion_planning import BruteForceFastMotionPlanner, HeuristicFastMotionPlanner, generate_component, generate_default_object

import numpy as np

from itertools import product, combinations
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import time

x_comp = {'id': 'x', 'pos': 0.0, 'pmin': 4.0, 'pmax': 4.0, 'vel':3.0, 'accel':1.0}
y_comp = {'id': 'y', 'pos': 1.5, 'pmin': 2.8, 'pmax': 3.2, 'vel':3.0, 'accel':0.5}
z_comp = {'id': 'z', 'pos': 1.0, 'pmin': 1.7, 'pmax': 2.3, 'vel':1.0, 'accel':0.1}

obstacles = np.loadtxt('tests/examples/obs10.csv', delimiter=',', skiprows=1)

# parameters
dt = 5
epsilon = 0.01
time_start = 0.0
num_solutions = None # Generate all values

# Format components and objects (Using list/ dictionary unpacking)
_x_comp = generate_component(**x_comp)
_y_comp = generate_component(**y_comp)
_z_comp = generate_component(**z_comp)
_obstacles = list(map(lambda x: generate_default_object(*x[:6]), obstacles))

now = time.time()

f = BruteForceFastMotionPlanner(dt, epsilon)
# f = HeuristicFastMotionPlanner(dt, epsilon) // If you want to use the heuristic method solver
sol = f.solve(time_start, _x_comp, _y_comp, _z_comp, _obstacles, n=num_solutions)

print(f"Found {len(sol)} solutions")
sample_rate = 100
num_traj_samples = 100
trajs = f.generate_trajectories_random_sample(time_start, sample_rate, num_traj_samples)

print(f"Found {len(trajs)} trajectories")
fig = plt.figure('3D Trajectory With Obstacles')
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')
ax.set_zlabel('Z axis')

# draw cuboid obstacles
for obs in obstacles:
    obsx = obs[0:2]
    obsy = obs[2:4]
    obsz = obs[4:6]
    for s, e in combinations(np.array(list(product(obsx, obsy, obsz))), 2):
        if np.count_nonzero(np.abs(s-e) == 0) == 2: # Only plot vertical or horizontal lines
            ax.plot3D(*zip(s, e), color="b")

for s, t in zip(sol, trajs):
    xpos = t['x']['pos']
    ypos = t['y']['pos']
    zpos = t['z']['pos']
    ax.plot(xpos, ypos, zpos)

fig, ax = plt.subplots(3, 1)
fig.canvas.set_window_title('Velocity')
ax[0].set_title('x vel')
ax[1].set_title('y vel')
ax[2].set_title('z vel')
for s, t in zip(sol, trajs):
    ax[0].plot(t['x']['vels'])
    ax[1].plot(t['y']['vels'])
    ax[2].plot(t['z']['vels'])

fig, ax = plt.subplots(3, 1)
fig.canvas.set_window_title('Acceleration')
ax[0].set_title('x accel')
ax[1].set_title('y accel')
ax[2].set_title('z accel')
for s, t in zip(sol, trajs):
    ax[0].plot(t['x']['accels'])
    ax[1].plot(t['y']['accels'])
    ax[2].plot(t['z']['accels'])

fig, ax = plt.subplots(3, 1)
fig.canvas.set_window_title('Jerk')
ax[0].set_title('x jerk')
ax[1].set_title('y jerk')
ax[2].set_title('z jerk')
for s, t in zip(sol, trajs):
    ax[0].plot(t['x']['jerks'])
    ax[1].plot(t['y']['jerks'])
    ax[2].plot(t['z']['jerks'])

plt.show()