from rrt_node import *
import matplotlib.pyplot as plt
from sympy import *
from sympy.geometry import *

def check_if_feasible(c1, c2, obstacles):
    for obstacle in obstacles:
        if len(obstacle.intersection(Segment(Point(c2[0,0], c2[1,0]), Point(c1[0,0], c1[1,0])))) > 0:
            return False
    return True


def sample_configuration():
    return np.random.uniform(np.array([0.0, 0.0]).reshape(2, 1), np.array([6.0, 6.0]).reshape(2, 1))


def extend_rrt(rrt_nodes, rejected_nodes, configuration, count, threshold, obstacles, heuristic=False):
    epsilon = 0.15
    best_rrt_node = rrt_nodes[0]
    best_rejected_node = rejected_nodes[0]

    for node in rejected_nodes:
        if np.linalg.norm(node.node_pos - configuration) < np.linalg.norm(best_rejected_node.node_pos - configuration):
            best_rejected_node = node

    for node in rrt_nodes:
        if np.linalg.norm(node.node_pos - configuration) < np.linalg.norm(best_rrt_node.node_pos - configuration):
            best_rrt_node = node

    new_rej_configuration = best_rejected_node.node_pos + epsilon * (
            configuration - best_rejected_node.node_pos) / np.linalg.norm(
        configuration - best_rejected_node.node_pos)

    if check_if_feasible(new_rej_configuration, best_rejected_node.node_pos, obstacles) and  np.linalg.norm(best_rejected_node.node_pos - rej_nodes[0].node_pos) > 0:
        rej_nodes.append(TreeNode(new_rej_configuration))

    if (heuristic is True and count > threshold and np.linalg.norm(best_rejected_node.node_pos - configuration) < np.linalg.norm(
            best_rrt_node.node_pos - configuration)):
        return 2

    new_configuration = best_rrt_node.node_pos + epsilon * (configuration - best_rrt_node.node_pos) / np.linalg.norm(configuration - best_rrt_node.node_pos)

    if heuristic is True:
        epsilon_counter = 0
        for node in nodes:
            if np.linalg.norm(node.node_pos - new_configuration) < 2.5 * epsilon:
                epsilon_counter += 1

        if epsilon_counter >= 5:
            return 2

    if not check_if_feasible(new_configuration, best_rrt_node.node_pos, obstacles):
        rej_nodes.append(TreeNode(new_configuration))
        return 3

    new_node = TreeNode(new_configuration, best_rrt_node)
    rrt_nodes.append(new_node)
    best_rrt_node.children.append(new_node)
    return 1


if __name__ == '__main__':
    epsilon = 0.15
    exp_nodes = []
    map_nodes_ = [np.array([float(i % int( 6 / epsilon)) * epsilon * 0.5, float(i / int(6 / epsilon)) * epsilon * 0.5]).reshape(2, 1) for i in range(int( 6 * 6 / (epsilon * epsilon)))]

    nodes = [TreeNode(np.array([1, 2]))]
    rej_nodes = [TreeNode(np.array([10, 10]).reshape(2, 1))]
    obstacles = [Segment(Point(0, 1.8), Point(5.5, 1.8)), Segment(Point(0, 2.2), Point(5.2, 2.2)), Segment(Point(0, 1.8), Point(0, 8)),
                 Segment(Point(5.2, 2.2), Point(5.2, 4)), Segment(Point(5.5, 1.8), Point(5.5, 4)), Segment(Point(0, 4), Point(5.2, 4)),
                 Segment(Point(5.5, 4), Point(6, 4))]

    """obstacles = [Segment(Point(0, 2), Point(2, 2)),
                 Segment(Point(2, 2), Point(2, 1.5)),
                 Segment(Point(2, 0), Point(2, 1)),
                 Segment(Point(2, 1), Point(6, 1)),
                 Segment(Point(2, 1.5), Point(5.5, 1.5)),
                 Segment(Point(5.5, 1.5), Point(5.5, 5.5)),
                 Segment(Point(6, 1), Point(6, 6)),
                 Segment(Point(5.5, 5.5), Point(1.5, 5.5)),
                 Segment(Point(6, 6), Point(1, 6)),
                 Segment(Point(1.5, 5.5), Point(1.5, 3)),
                 Segment(Point(1, 6), Point(1, 2.5)),
                 Segment(Point(1.5, 3), Point(4.5, 3)),
                 Segment(Point(1, 2.5), Point(5, 2.5)),
                 Segment(Point(4.5, 3), Point(4.5, 4.5)),
                 Segment(Point(5, 2.5), Point(5, 5)),
                 Segment(Point(4.5, 4.5), Point(2.5, 4.5)),
                 Segment(Point(5, 5), Point(2, 5))]"""

    fig = plt.figure()
    plt.ylim(0, 6)
    plt.xlim(0, 6)

    i = 0
    cnt = 0
    while i < 1000:
        conf = sample_configuration()
        z = extend_rrt(nodes, rej_nodes, conf, cnt, 100, obstacles, heuristic=True)
        if z != 2:
            exp_nodes.append(conf)
            cnt += 1
        elif z == 2:
            i -= 1
        i += 1                                                      #3.5 3.5
        if np.linalg.norm(nodes[len(nodes) - 1].node_pos - np.array([2.5, 5.1]).reshape(2, 1)) < 0.25:
            print('break')
            break


    print(len(rej_nodes), cnt)
    input()

    for obstacle in obstacles:
        x1, y1 = [obstacle.p1[0], obstacle.p2[0]], [obstacle.p1[1], obstacle.p2[1]]
        plt.plot(x1, y1, marker='.', color='r', markersize=0.5)


    rej_nodes.pop(0)
    for node_ in nodes:
        if node_.parent_node is not None:
            x1, y1 = [node_.node_pos[0, 0], node_.parent_node.node_pos[0, 0]], [node_.node_pos[1, 0], node_.parent_node.node_pos[1, 0]]
            plt.plot(x1, y1, marker='.', color='b', markersize=0.1)
            fig.canvas.draw()
            plt.pause(0.001)

    for node_ in rej_nodes:
        plt.scatter(node_.node_pos[0, 0], node_.node_pos[1, 0], s=4, color='r')
        fig.canvas.draw()
        plt.pause(0.001)

    plt.show()