from collections import namedtuple
from fast_motion_planning import BruteForceFastMotionPlanner, 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 = [
    [3.75, 4.25, 2.3, 2.8, 1.7, 2.3],
    [3.75, 4.25, 3.2, 3.7, 1.7, 2.3]
]

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

fig = plt.figure()
# plt.ion()
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")

# parameters
dt = 5
epsilon = 0.01
time_start = 0.0
sample_rate = 10
num_traj_samples = 200

# 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))

f = BruteForceFastMotionPlanner(dt, epsilon)
solution_trajectories_iterator = f.solve_trajectories_iter(
    time_start, _x_comp, _y_comp, _z_comp, _obstacles, sample_rate)

for i, (sol, traj) in enumerate(solution_trajectories_iterator):
    xpos = traj['x']['pos']
    ypos = traj['y']['pos']
    zpos = traj['z']['pos']
    ax.plot(xpos, ypos, zpos)

    # plt.draw()
    # plt.pause(0.001)
    # # plt.clf()
    if i > num_traj_samples:
        break


plt.show(block=True)