"""
This is rrt star code for 3D
@author: yue qi
"""
import numpy as np
from numpy.matlib import repmat
from rrt_3D.env3D import env
from collections import defaultdict
import pyrr as pyrr
from utils3D import getDist, sampleFree, nearest, steer, isCollide, near, visualization, cost, path
import time


class rrtstar():
    def __init__(self):
        self.env = env()
        self.Parent = defaultdict(lambda: defaultdict(dict))
        self.V = []
        self.E = []
        self.i = 0
        self.maxiter = 10000
        self.stepsize = 0.5
        self.Path = []

    def wireup(self,x,y):
        self.E.append([x,y]) # add edge
        self.Parent[str(x[0])][str(x[1])][str(x[2])] = y

    def removewire(self,xnear):
        xparent = self.Parent[str(xnear[0])][str(xnear[1])][str(xnear[2])]
        a = np.array([xnear,xparent])
        self.E = [xx for xx in self.E if not (xx==a).all()] # remove and replace old the connection

    def run(self):
        self.V.append(self.env.start)
        ind = 0
        xnew = self.env.start
        while ind < self.maxiter and getDist(xnew,self.env.goal) > 1:
            xrand    = sampleFree(self)
            xnearest = nearest(self,xrand)
            xnew     = steer(self,xnearest,xrand)
            if not isCollide(self,xnearest,xnew):
                self.V.append(xnew) # add point
                self.wireup(xnew,xnearest)
                #visualization(self)
                self.i += 1
            ind += 1
            if getDist(xnew,self.env.goal) <= 1:
                self.wireup(self.env.goal,xnew)
                self.Path,D = path(self)
                print('Total distance = '+str(D))
        visualization(self)

if __name__ == '__main__':
    p = rrtstar()
    starttime = time.time()
    p.run()
    print('time used = ' + str(time.time()-starttime))