import numpy as np
from numpy.matlib import repmat
import pyrr as pyrr
# plotting
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import mpl_toolkits.mplot3d as plt3d


def getRay(x, y):
    direc = [y[0] - x[0], y[1] - x[1], y[2] - x[2]]
    return np.array([x, direc])


def getAABB(blocks):
    AABB = []
    for i in blocks:
        AABB.append(np.array([np.add(i[0:3], -0), np.add(i[3:6], 0)]))  # make AABBs alittle bit of larger
    return AABB


def getDist(pos1, pos2):
    return np.sqrt(sum([(pos1[0] - pos2[0]) ** 2, (pos1[1] - pos2[1]) ** 2, (pos1[2] - pos2[2]) ** 2]))


def draw_block_list(ax, blocks):
    '''
    Subroutine used by draw_map() to display the environment blocks
    '''
    v = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1]],
                 dtype='float')
    f = np.array([[0, 1, 5, 4], [1, 2, 6, 5], [2, 3, 7, 6], [3, 0, 4, 7], [0, 1, 2, 3], [4, 5, 6, 7]])
    # clr = blocks[:,6:]/255
    n = blocks.shape[0]
    d = blocks[:, 3:6] - blocks[:, :3]
    vl = np.zeros((8 * n, 3))
    fl = np.zeros((6 * n, 4), dtype='int64')
    # fcl = np.zeros((6*n,3))
    for k in range(n):
        vl[k * 8:(k + 1) * 8, :] = v * d[k] + blocks[k, :3]
        fl[k * 6:(k + 1) * 6, :] = f + k * 8
        # fcl[k*6:(k+1)*6,:] = clr[k,:]

    if type(ax) is Poly3DCollection:
        ax.set_verts(vl[fl])
    else:
        pc = Poly3DCollection(vl[fl], alpha=0.15, linewidths=1, edgecolors='k')
        # pc.set_facecolor(fcl)
        h = ax.add_collection3d(pc)
        return h


''' The following utils can be used for rrt or rrt*,
    required param initparams should have
    env,      environement generated from env3D
    V,        node set
    E,        edge set
    i,        nodes added
    maxiter,  maximum iteration allowed
    stepsize, leaf growth restriction

'''


def sampleFree(initparams):
    x = np.random.uniform(initparams.env.boundary[0:3], initparams.env.boundary[3:6])
    if isinside(initparams, x):
        return sampleFree(initparams)
    else:
        return np.array(x)


def isinside(initparams, x):
    '''see if inside obstacle'''
    for i in initparams.env.blocks:
        if i[0] <= x[0] < i[3] and i[1] <= x[1] < i[4] and i[2] <= x[2] < i[5]:
            return True
    return False


def isCollide(initparams, x, y):
    '''see if line intersects obstacle'''
    ray = getRay(x, y)
    dist = getDist(x, y)
    for i in getAABB(initparams.env.blocks):
        shot = pyrr.geometric_tests.ray_intersect_aabb(ray, i)
        if shot is not None:
            dist_wall = getDist(x, shot)
            if dist_wall <= dist:  # collide
                return True
    return False


def nearest(initparams, x):
    V = np.array(initparams.V)
    if initparams.i == 0:
        return initparams.V[0]
    xr = repmat(x, len(V), 1)
    dists = np.linalg.norm(xr - V, axis=1)
    return initparams.V[np.argmin(dists)]


def steer(initparams, x, y):
    direc = (y - x) / np.linalg.norm(y - x)
    xnew = x + initparams.stepsize * direc
    return xnew


def near(initparams, x, r=2):
    # TODO: r = min{gamma*log(card(V)/card(V)1/d),eta}
    V = np.array(initparams.V)
    if initparams.i == 0:
        return initparams.V[0]
    xr = repmat(x, len(V), 1)
    inside = np.linalg.norm(xr - V, axis=1) < r
    nearpoints = V[inside]
    return np.array(nearpoints)


def cost(initparams, x):
    '''here use the additive recursive cost function'''
    if all(x == initparams.env.start):
        return 0
    xparent = initparams.Parent[str(x[0])][str(x[1])][str(x[2])]
    return cost(initparams, xparent) + getDist(x, xparent)


def visualization(initparams):
    V = np.array(initparams.V)
    E = np.array(initparams.E)
    Path = np.array(initparams.Path)
    start = initparams.env.start
    goal = initparams.env.goal
    ax = plt.subplot(111, projection='3d')
    ax.view_init(elev=0., azim=90)
    ax.clear()
    draw_block_list(ax, initparams.env.blocks)
    if E != []:
        for i in E:
            xs = i[0][0], i[1][0]
            ys = i[0][1], i[1][1]
            zs = i[0][2], i[1][2]
            line = plt3d.art3d.Line3D(xs, ys, zs)
            ax.add_line(line)

    if Path != []:
        for i in Path:
            xs = i[0][0], i[1][0]
            ys = i[0][1], i[1][1]
            zs = i[0][2], i[1][2]
            line = plt3d.art3d.Line3D(xs, ys, zs, color='r')
            ax.add_line(line)

    ax.plot(start[0:1], start[1:2], start[2:], 'go', markersize=7, markeredgecolor='k')
    ax.plot(goal[0:1], goal[1:2], goal[2:], 'ro', markersize=7, markeredgecolor='k')
    ax.scatter3D(V[:, 0], V[:, 1], V[:, 2])
    plt.xlim(initparams.env.boundary[0], initparams.env.boundary[3])
    plt.ylim(initparams.env.boundary[1], initparams.env.boundary[4])
    ax.set_zlim(initparams.env.boundary[2], initparams.env.boundary[5])
    plt.xlabel('x')
    plt.ylabel('y')
    if not Path != []:
        plt.pause(0.001)
    else:
        plt.show()


def path(initparams, Path=[], dist=0):
    x = initparams.env.goal
    while not all(x == initparams.env.start):
        x2 = initparams.Parent[str(x[0])][str(x[1])][str(x[2])]
        Path.append(np.array([x, x2]))
        dist += getDist(x, x2)
        x = x2
    return Path, dist