# AUTOGENERATED! DO NOT EDIT! File to edit: driving.ipynb (unless otherwise specified).

__all__ = ['BG_COLOR', 'GRID_COLOR', 'read_ply', 'sys_byteorder', 'ply_dtypes', 'read_lidar_points', 'cloud_layout',
           'visualize_cloud', 'gen_color_palette', 'COLOR_OPTIONS', 'visualize_clouds', 'animation_layout',
           'visualize_clouds_animation']

# Cell
import sys
import numpy as np
import pandas as pd
from collections import defaultdict

import plotly.express as px
import plotly.graph_objects as go

BG_COLOR = 'rgb(30, 30, 30)'
GRID_COLOR = 'rgb(127, 127, 127)'


# Cell
# Adapted from pyntcloud under MIT license

sys_byteorder = ('>', '<')[sys.byteorder == 'little']

ply_dtypes = dict([
    (b'int8', 'i1'),
    (b'char', 'i1'),
    (b'uint8', 'u1'),
    (b'uchar', 'b1'),
    (b'uchar', 'u1'),
    (b'int16', 'i2'),
    (b'short', 'i2'),
    (b'uint16', 'u2'),
    (b'ushort', 'u2'),
    (b'int32', 'i4'),
    (b'int', 'i4'),
    (b'uint32', 'u4'),
    (b'uint', 'u4'),
    (b'float32', 'f4'),
    (b'float', 'f4'),
    (b'float64', 'f8'),
    (b'double', 'f8')
])


def read_ply(filename):
    """ Read a binary_little_endian .ply file and return data as a dict.

    Parameters:
        filename: of ply file

    Returns:
        A dictionary with `points`, `mesh`, and/or `comments` keys.
    """
    with open(filename, 'rb') as ply:

        if b'ply' not in ply.readline():
            raise ValueError('The file does not start with the word ply')

        # make sure format is binary_little_endian
        fmt = ply.readline().split()[1].decode()
        assert fmt == 'binary_little_endian'

        line = []
        dtypes = defaultdict(list)
        count = 2
        points_size = None
        mesh_size = None
        comments = []
        while b'end_header' not in line and line != b'':
            line = ply.readline()

            if b'element' in line:
                line = line.split()
                name = line[1].decode()
                size = int(line[2])
                if name == "vertex":
                    points_size = size
                elif name == "face":
                    mesh_size = size

            elif b'property' in line:
                line = line.split()
                # element mesh
                if b'list' in line:

                    if b"vertex_indices" in line[-1] or b"vertex_index" in line[-1]:
                        mesh_names = ["n_points", "v1", "v2", "v3"]
                    else:
                        mesh_names = ["n_coords"] + ["v1_u", "v1_v", "v2_u",
                                                     "v2_v", "v3_u", "v3_v"]

                    # the first number has different dtype than the list
                    dtypes[name].append(
                        (mesh_names[0], '<' + ply_dtypes[line[2]]))
                    # rest of the numbers have the same dtype
                    dt = '<' + ply_dtypes[line[3]]

                    for j in range(1, len(mesh_names)):
                        dtypes[name].append((mesh_names[j], dt))
                else:
                    dtypes[name].append(
                        (line[2].decode(), '<' + ply_dtypes[line[1]]))

            elif b'comment' in line:
                line = line.split(b" ", 1)
                comment = line[1].decode().rstrip()
                comments.append(comment)

            count += 1

        end_header = ply.tell()

    data = {}

    if comments:
        data["comments"] = comments

    with open(filename, 'rb') as ply:
        ply.seek(end_header)
        points_np = np.fromfile(ply, dtype=dtypes["vertex"], count=points_size)
        if sys_byteorder != '<':
            points_np = points_np.byteswap().newbyteorder()
        data["points"] = pd.DataFrame(points_np)
        if mesh_size:
            mesh_np = np.fromfile(ply, dtype=dtypes["face"], count=mesh_size)
            if sys_byteorder != '<':
                mesh_np = mesh_np.byteswap().newbyteorder()
            data["mesh"] = pd.DataFrame(mesh_np)
            data["mesh"].drop('n_points', axis=1, inplace=True)

    return data


# Cell
# adapted from code by 3630 TAs Binit Shah and Jerred Chen
def read_lidar_points(filename):
    """ Read 3D points in LIDAR scan stored as a binary_little_endian .ply file.

    Parameters:
        filename: of ply file

    Returns:
        A tuple (3,N) numpy array.
    """
    data = read_ply(filename)
    points = data["points"]
    np_cloud = np.empty((3, len(points)))
    np_cloud[0, :] = points['x']
    np_cloud[1, :] = points['y']
    np_cloud[2, :] = points['z']
    return np_cloud


# Cell
def cloud_layout(show_grid_lines):
    """Create layout for showing clouds."""
    grid_lines_color = GRID_COLOR if show_grid_lines else BG_COLOR
    layout = go.Layout(
        scene=dict(
            xaxis=dict(nticks=8,
                       showbackground=True,
                       backgroundcolor=BG_COLOR,
                       gridcolor=grid_lines_color,
                       zerolinecolor=grid_lines_color),
            yaxis=dict(nticks=8,
                       showbackground=True,
                       backgroundcolor=BG_COLOR,
                       gridcolor=grid_lines_color,
                       zerolinecolor=grid_lines_color),
            zaxis=dict(nticks=8,
                       showbackground=True,
                       backgroundcolor=BG_COLOR,
                       gridcolor=grid_lines_color,
                       zerolinecolor=grid_lines_color),
            xaxis_title="x (meters)",
            yaxis_title="y (meters)",
            zaxis_title="z (meters)"
        ),
        scene_aspectmode='data',
        margin=dict(r=10, l=10, b=10, t=10),
        paper_bgcolor=BG_COLOR,
        font=dict(
            family="Courier New, monospace",
            color=grid_lines_color
        ),
        legend=dict(
            font=dict(
                family="Courier New, monospace",
                color=GRID_COLOR
            )
        )
    )
    return layout


# Cell
def visualize_cloud(cloud, show_grid_lines=False,  color='#90FF90',
                    marker_size=1, fraction=None):
    """ Visualizes point cloud in 3D scatter plot.

    Args:
        cloud (np.ndarray):     point cloud, a (3, num_points) numpy array
        show_grid_lines (bool): plots gridlines
        color (str):            color for markers
        marker_size (int):      size of each marker
        fraction (double):      take only a fraction of the points
    """
    # Setup data
    data = []
    N = cloud.shape[1]
    if fraction is not None:
        subset = np.random.choice(N, int(N * fraction), replace=False)
    x_data = cloud[0][subset] if fraction is not None else cloud[0]
    y_data = cloud[1][subset] if fraction is not None else cloud[1]
    z_data = cloud[2][subset] if fraction is not None else cloud[2]

    data.append(go.Scatter3d(x=x_data, y=y_data, z=z_data,
                             mode='markers',
                             marker=dict(
                                 size=marker_size,
                                 color=color,
                                 opacity=1.0
                             ))
                )
    fig = go.Figure(data=data, layout=cloud_layout(show_grid_lines))
    fig.show()


# Cell
COLOR_OPTIONS = ['#FF1493', '#CD5C5C', '#FFDAB9', '#FF4500', '#8B0000', '#E6E6FA', '#87CEFA', '#CD853F', '#DC143C',
                 '#808000', '#483D8B', '#D2691E', '#FF69B4', '#8FBC8F', '#A0522D', '#9932CC', '#B22222', '#F08080',
                 '#FA8072', '#4B0082', '#EE82EE', '#008000', '#F0F8FF', '#00FA9A', '#FFA500', '#BA55D3', '#0000FF',
                 '#66CDAA', '#FFF8DC', '#9370DB', '#00BFFF', '#FFFFFF', '#FF0000', '#3CB371', '#4682B4', '#FFFAF0',
                 '#FFFFE0', '#FFD700', '#800000', '#A52A2A', '#7FFF00', '#DDA0DD', '#F5F5F5', '#EEE8AA', '#F5F5DC',
                 '#FFC0CB', '#6495ED', '#8A2BE2', '#C0C0C0', '#F5FFFA', '#FF8C00', '#20B2AA', '#48D1CC', '#E0FFFF',
                 '#87CEEB', '#4169E1', '#FF7F50', '#F8F8FF', '#DAA520', '#B0C4DE', '#F4A460', '#00CED1', '#2E8B57',
                 '#7CFC00', '#7FFFD4', '#FFB6C1', '#B8860B', '#8B4513', '#8B008B', '#BC8F8F', '#663399', '#C71585',
                 '#F0FFF0', '#D2B48C', '#F0E68C', '#00FF00', '#BDB76B', '#5F9EA0', '#ADD8E6', '#F0FFFF', '#1E90FF',
                 '#FFF5EE', '#FFA07A', '#778899', '#ADFF2F', '#DB7093', '#FFE4E1', '#FF00FF', '#32CD32', '#FFFF00',
                 '#DCDCDC', '#9ACD32', '#FFDEAD', '#DA70D6', '#008B8B', '#6A5ACD', '#008080', '#D8BFD8', '#00FF7F',
                 '#FF6347', '#228B22', '#6B8E23', '#708090', '#556B2F', '#40E0D0', '#98FB98', '#90EE90', '#7B68EE',
                 '#696969', '#E9967A', '#00FFFF', '#F5DEB3', '#FFFACD', '#D3D3D3', '#AFEEEE', '#FFF0F5', '#191970']


def gen_color_palette(n):
    """Generates a hex color palette of size n, without repeats
    and only light colors (easily visible on dark background).

    Args:
        n (int): number of clouds, each cloud gets a unique color
    """
    palette = []
    do_replace = False if len(COLOR_OPTIONS) >= n else True
    for i in np.random.choice(len(COLOR_OPTIONS), n, replace=do_replace):
        palette.append(COLOR_OPTIONS[i])

    return palette


# Cell
def visualize_clouds(clouds, show_grid_lines=False, cloud_colors=None,
                     marker_size=1, do_subsampling=True):
    """Visualizes cloud(s) in a iterative 3D plot.
    Due to browser limitations, rendering above 5 frames requires
    subsampling of the point clouds, which is done automatically.

    Example input of arg:
    clouds = [clouda, cloudb, cloudc]
    where each cloud is a numpy array of shape (3, num_points).
    cloud[0] are the x coordinates, cloud[1] is y, and cloud[2] is z.

    Args:
        clouds (list):          ordered series of point clouds
        show_grid_lines (bool): plots gridlines
        cloud_colors (list):    colors for each cloud in the visualization
        marker_size (int):      size of each marker
        do_subsampling (bool):  whether or not subsampling occurs
    """
    # Setup data
    nc = len(clouds)
    palette = gen_color_palette(nc)
    if cloud_colors is not None:
        if isinstance(cloud_colors, list):
            if nc != len(cloud_colors):
                raise ValueError(
                    'length of cloud_colors does not match length of clouds')
        else:
            cloud_colors = [cloud_colors] * nc
    data = []
    for i, cloud in enumerate(clouds):
        N = cloud.shape[1]
        subset = np.random.choice(N,
                                  int(N / (nc * 0.25)) if do_subsampling and int(
                                      N / (nc * 0.25)) <= N else N,
                                  replace=False)
        x_data = cloud[0][subset] if nc > 5 else cloud[0]
        y_data = cloud[1][subset] if nc > 5 else cloud[1]
        z_data = cloud[2][subset] if nc > 5 else cloud[2]

        data.append(go.Scatter3d(x=x_data, y=y_data, z=z_data,
                                 mode='markers',
                                 marker=dict(
                                     size=marker_size,
                                     color=palette[i] if not cloud_colors else cloud_colors[i],
                                     opacity=1.0
                                 ))
                    )

    fig = go.Figure(data=data, layout=cloud_layout(show_grid_lines))
    fig.show()


# Cell
def animation_layout(clouds_labels, speed):
    """Setup layout for animation."""
    return go.Layout(
        hovermode=False,
        sliders=[dict(
            active=0,
            yanchor="top",
            xanchor="left",
            currentvalue=dict(
                font=dict(
                    family="Courier New, monospace",
                    color=BG_COLOR
                ),
                prefix="Step ",
                visible=True,
                xanchor="right"
            ),
            pad=dict(b=10, t=0),
            len=0.95,
            x=0.05,
            y=0,
            font=dict(
                family="Courier New, monospace",
                color=GRID_COLOR
            ),
            steps=[dict(
                args=[[clouds_label], dict(
                    frame=dict(duration=speed),
                    mode="immediate"
                )],
                label=clouds_label,
                method="animate") for clouds_label in clouds_labels]
        )],
        updatemenus=[dict(
            buttons=[dict(
                args=[None, dict(
                    frame=dict(duration=speed),
                    fromcurrent=True
                )],
                label="Play",
                method="animate"
            )],
            direction="left",
            pad=dict(r=10, t=30),
            showactive=False,
            type="buttons",
            x=0.05,
            xanchor="right",
            y=0,
            yanchor="top",
            font=dict(
                family="Courier New, monospace",
                color=GRID_COLOR
            )
        )],
    )


# Cell
def visualize_clouds_animation(clouds_series, clouds_labels=None,
                               show_grid_lines=False, speed=100,
                               cloud_colors=None, marker_size=1,
                               do_subsampling=True):
    """Visualizes a cloud(s) series using a slider & play/pause.
    Due to browser limitations, rendering above 5 frames requires
    subsampling of the point clouds, which is done automatically.

    Example input of arg:
    clouds_series = [[clouda, cloudb], [clouda, cloudb]]
    clouds_labels = ["iteration1", "iteration2"]

    Args:
        clouds_series (list):   ordered series of point clouds
        clouds_labels (list):   ordered labels for clouds
        show_grid_lines (bool): plots gridlines
        speed (int):            speed at which the frames are played through
        cloud_colors (list):    colors for each cloud in the visualization
        marker_size (int):      size of each marker
        do_subsampling (bool):  whether or not subsampling occurs
    """
    clouds_labels = [str(i) for i in range(len(clouds_series))
                     ] if not clouds_labels else clouds_labels

    first_data = None
    palette = gen_color_palette(len(clouds_series[0]))
    if cloud_colors is not None:
        if isinstance(cloud_colors, list):
            if len(clouds_series[0]) != len(cloud_colors):
                raise ValueError(
                    'length of cloud_colors does not match length of clouds')
        else:
            raise ValueError('cloud_colors must be a list')
    frames = []
    for i, (label, clouds) in enumerate(zip(clouds_labels, clouds_series)):
        data = []
        for j, cloud in enumerate(clouds):
            nc = len(clouds)
            N = cloud.shape[1]
            subset = np.random.choice(N,
                                      int(N / (nc * 0.25)) if do_subsampling and int(
                                          N / (nc * 0.25)) <= N else N,
                                      replace=False)
            x_data = cloud[0][subset] if nc > 5 else cloud[0]
            y_data = cloud[1][subset] if nc > 5 else cloud[1]
            z_data = cloud[2][subset] if nc > 5 else cloud[2]

            data.append(go.Scatter3d(x=x_data, y=y_data, z=z_data,
                                     mode='markers',
                                     marker=dict(
                                         size=marker_size,
                                         color=palette[j] if not cloud_colors else cloud_colors[j],
                                         opacity=1.0
                                     ))
                        )
        frames.append(go.Frame(name=label, data=data))
        if first_data is None:
            first_data = data

    fig = go.Figure(data=first_data, layout=cloud_layout(
        show_grid_lines), frames=frames)
    fig.update_layout(animation_layout(clouds_labels, speed))
    fig.show()
