#!/usr/bin/env python
"""Basic implementation of RRT for the MM."""

import bisect
import numpy as np
import matplotlib.pyplot as plt
import IPython

from mm2d import models

import geometry as geo


def collision_robot(model, obstacles, q):
    """Check for collision between the robot and a circle shape."""
    # see http://jeffreythompson.org/collision-detection/circle-rect.php
    # construct AABB out of external points on the base + two on arm
    ax, ay = model.arm_points(q)
    bx, by = model.base_corners(q)
    x = np.concatenate((ax, bx))
    y = np.concatenate((ay, by))
    points = np.vstack((x, y)).T
    aabb_robot = geo.aabb(points)

    for obstacle in obstacles:
        if type(obstacle) == geo.Floor:
            if (ay < obstacle.y).any():
                return True
        elif type(obstacle) == geo.Circle:
            circle = obstacle
            if not geo.collision_rect_circle(aabb_robot, circle):
                continue

            # check collision with base
            base_rect = geo.aabb(np.vstack((bx, by)).T)
            if geo.collision_rect_circle(base_rect, circle):
                return True

            # check collision with link 1
            _, link1_dist2 = geo.dist_to_edge(
                np.array([ax[0], ay[0]]), np.array([ax[1], ay[1]]), circle.center
            )
            if link1_dist2 < circle.r ** 2:
                return True

            # check collision with link 2
            _, link2_dist2 = geo.dist_to_edge(
                np.array([ax[1], ay[1]]), np.array([ax[2], ay[2]]), circle.center
            )
            if link2_dist2 < circle.r ** 2:
                return True

    return False


def dist_squared(q1, q2):
    """Squared distance between two configurations."""
    q = q2 - q1
    # TODO not sure if I have to handle wrapping: with joint limits it may not
    # make sense
    # q[1:] = q[1:] % (2 * np.pi)
    return q @ q


def dist(q1, q2):
    return np.sqrt(dist_squared(q1, q2))


def very_close(q1, q2, eps=1e-8):
    """True if points q1 and q2 are very close together."""
    return dist_squared(q1, q2) < eps ** 2


class Segment:
    def __init__(self, v1, v2):
        self.v1 = v1
        self.v2 = v2
        self.length = v2.cost
        self.divs = [(v1, 0), (v2, v2.cost)]

    def divide(self, vertex):
        """Add a division to the segment."""
        d = dist(vertex.value, self.v1.value)
        self.divs.append((vertex, d))
        self.divs.sort(key=lambda div: div[1])  # keep sorted

    def get_edge(self, vertex):
        """Return end points of edge containing dist."""
        d = dist(vertex.value, self.v1.value)
        dists = [div[1] for div in self.divs]
        i = bisect.bisect(dists, d)
        return self.divs[i - 1][0], self.divs[i][0]

    def ends(self):
        return self.v1, self.v2


class RRTNode:
    def __init__(self, value, idx):
        self.value = value
        self.idx = idx

        self.parent_idx = -1
        self.children_idx = set()
        self.cost = np.infty

    def add_child(self, other):
        """Add a child node.

        This changes both this node and the parent idx of the other node.
        """
        other.parent_idx = self.idx
        other.cost = dist(self.value, other.value)
        if other.cost < 1e-8:
            print(f"0 cost: {self.idx}, {other.idx}")
        self.children_idx.add(other.idx)

    def remove_child(self, other):
        """Remove a child node."""
        self.children_idx.remove(other.idx)


class RRTree:
    """Rapidly-exploring random tree."""

    def __init__(self, model, root, obstacles=[]):
        self.model = model
        self.vertices = [RRTNode(root, 0)]

        # Segments are straight lines connecting two vertices.
        self.segments = []

        self.obstacles = obstacles

    def _nearest(self, q):
        """Get nearest point in the tree to the given value. This can be either
        an existing vertex or an edge."""
        min_dist_idx = -1
        min_dist = np.infty
        q_closest = np.zeros(3)
        for idx, segment in enumerate(self.segments):
            q1 = segment.v1.value
            q2 = segment.v2.value
            qc, dist2 = geo.dist_to_edge(q1, q2, q)
            if dist2 < min_dist:
                min_dist_idx = idx
                min_dist = dist2
                q_closest = qc
        return min_dist_idx, q_closest

    def _stop_config(self, q_near, q_samp):
        """Get the stopping configuration closest to q_samp.

        Get the configuration closest to q_samp along the straight-line path
        from q_near that is obstacle free.
        """
        Δd = 0.02
        d = dist(q_near, q_samp)
        n = int(d / Δd) + 1  # number of samples
        ts = np.linspace(0, 1, n)

        for idx, t in enumerate(ts):
            q = (1 - t) * q_near + t * q_samp  # linear interp
            # for obstacle in self.obstacles:
            # if this configuration is in collision, return the last one
            # (which was collision free)
            if collision_robot(self.model, self.obstacles, q):
                t_last = ts[idx - 1]
                return (1 - t_last) * q_near + t_last * q_samp

        # the whole path is collision free, so we can add q_samp directly
        return q_samp

    def _cut_edge(self, v1, v2, new_vertex):
        if v2.parent_idx == v1.idx:
            parent = v1
            child = v2
        elif v1.parent_idx == v2.idx:
            parent = v2
            child = v1
        else:
            IPython.embed()
            raise Exception("Nodes are not connected")

        parent.remove_child(child)
        parent.add_child(new_vertex)
        new_vertex.add_child(child)

    def _add_segment(self, v1, v2):
        self.segments.append(Segment(v1, v2))

    def _add_first_vertex(self, q):
        """Add first vertex.

        This is a special case as there are no segments yet.
        """
        q0 = self.vertices[0].value
        q = self._stop_config(q0, q)
        if very_close(q0, q):
            return -1

        vertex = RRTNode(q, 1)
        self.vertices.append(vertex)
        self.vertices[0].add_child(vertex)

        # length = dist(q0, q)
        self.segments.append(Segment(self.vertices[0], vertex))

        return 1

    def add_vertex(self, q):
        """Add vertex and return its index."""
        # first addition
        if len(self.segments) == 0:
            return self._add_first_vertex(q)

        segment_idx, q_closest = self._nearest(q)
        segment = self.segments[segment_idx]

        q = self._stop_config(q_closest, q)
        if very_close(q, q_closest):
            return -1

        q_idx = len(self.vertices)

        v1, v2 = segment.ends()

        if very_close(v1.value, q_closest):
            new_vertex = RRTNode(q, q_idx)
            self.vertices.append(new_vertex)
            v1.add_child(new_vertex)
            self._add_segment(v1, new_vertex)

        elif very_close(v2.value, q_closest):
            new_vertex = RRTNode(q, q_idx)
            self.vertices.append(new_vertex)
            v2.add_child(new_vertex)
            self._add_segment(v2, new_vertex)

        else:
            # add new vertex
            new_vertex1 = RRTNode(q_closest, q_idx)
            self.vertices.append(new_vertex1)

            # mark the new division on the segment
            v1, v2 = segment.get_edge(new_vertex1)
            segment.divide(new_vertex1)

            # cut the edge between the old nodes and insert new vertex
            self._cut_edge(v1, v2, new_vertex1)

            q_idx += 1

            # add the edge to the new vertex q
            new_vertex2 = RRTNode(q, q_idx)
            self.vertices.append(new_vertex2)
            new_vertex1.add_child(new_vertex2)

            self._add_segment(new_vertex1, new_vertex2)

        return q_idx

    def path_to_root(self, idx):
        """Get path to root from the given idx.

        Returns a list of indices.
        """
        indices = [idx]
        vertex = self.vertices[idx]
        while vertex.parent_idx > -1:
            indices.append(vertex.parent_idx)
            vertex = self.vertices[vertex.parent_idx]
        return indices


def sample_config_space(low, high):
    """Generate a uniform random sample from configuration space."""
    assert low.shape == high.shape
    return np.random.random(low.shape[0]) * (high - low) + low


def main():
    np.set_printoptions(precision=3, suppress=True)
    np.random.seed(0)

    model = models.ThreeInputModel(output_idx=[0, 1])

    circle1 = geo.Circle(3, 1, 0.5)
    floor = geo.Floor(0)

    q_low = np.array([-1, -np.pi, -np.pi])
    q_high = np.array([6, np.pi, np.pi])
    q0 = np.array([0, np.pi / 4.0, -np.pi / 4.0])
    qf = np.array([5, np.pi / 4.0, -np.pi / 4.0])

    # q_low = np.array([-10, -10])
    # q_high = np.array([10, 10])
    # q0 = np.zeros(2)

    # IPython.embed()

    tree = RRTree(model, q0, obstacles=[circle1, floor])

    i = 0
    while True:
        # try to connect the final node every 100th sample
        if i % 100 == 0:
            q = qf
        else:
            q = sample_config_space(q_low, q_high)
        q_idx = tree.add_vertex(q)
        if q_idx >= 0:
            if very_close(tree.vertices[q_idx].value, qf):
                print("qf connected")
                break
        i += 1

    print(f"Number of vertices = {len(tree.vertices)}")
    print(f"Number of segments = {len(tree.segments)}")

    path = tree.path_to_root(q_idx)
    q_path = np.array([tree.vertices[idx].value for idx in path])
    print(f"Length of path = {len(path)}")

    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    for segment in tree.segments:
        q1 = segment.v1.value
        q2 = segment.v2.value
        ax.plot([q1[0], q2[0]], [q1[1], q2[1]], zs=[q1[2], q2[2]], color="k")

    for i, vertex in enumerate(tree.vertices):
        q = vertex.value
        ax.plot(q[0], q[1], ".", zs=q[2], color="r")

    ax.plot(q_path[:, 0], q_path[:, 1], zs=q_path[:, 2], color="b")

    points = np.array([model.sample_points(q)[2:, :] for q in q_path])

    plt.figure()
    ax = plt.gca()
    ax.set_aspect("equal")
    plt.plot(points[:, 0, 0], points[:, 0, 1], "o-", label="p0")
    plt.plot(points[:, 1, 0], points[:, 1, 1], "o-", label="p1")
    plt.plot(points[:, 2, 0], points[:, 2, 1], "o-", label="p2")

    ax = plt.gca()
    ax.add_patch(plt.Circle(circle1.center, circle1.r, color="k", fill=False))

    plt.legend()

    plt.grid()

    plt.show()


if __name__ == "__main__":
    main()
