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

__all__ = ['plot_trajectory', 'tortuosity', 'map_size', 'avg_distance_between_pair', 'random_color', 'draw_pair_points',
           'plot_all_env']

# Cell
import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
import pylab as plot
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import yaml
import math
import os
import random
from queue import Queue, PriorityQueue
plt.rcParams["figure.figsize"] = (16, 16)

# Cell
def plot_trajectory(data_root, env, quantity=None):
    """
        Calculate metrics for the input map.
        Args:
            data_root (String) : The path of folder env_exp_files. Default at 'Data/env_exp_files' .
            env (String) : Which environment was used.
            quantity(int) : The draw quantity of trajectories. Default is None means total.

        Returns:
            A figure (trajectory.png) would be saved.
    """
    # Default path & file name
    folder_map = 'npy' #The relative path to find the yaml file `obstacle_map.map`.
    map_file = 'obstacle_map.npy' #The file name of the 2D occupency grid map.
    folder_path = 'trajectory' #The relative path to find the yaml file of trajectories.
    folder_info = 'environment_info' #The relative path to find the yaml file of information about 2D occupency grid map.
    info_file = 'grid_map_info.yaml' #The file name of the yaml file about 2D occupency grid map.

    yaml_path = data_root + '/' + env + '/' + folder_info + '/' + info_file
    with open(yaml_path, 'r') as f:
            info_yaml_data = yaml.load(f, Loader=yaml.CLoader)

    map_path = data_root + '/' + env + '/' + folder_map + '/' + map_file
    vis_trajectory_map = np.load(map_path)

    trajectory_folder = data_root + '/' + env + '/' + folder_path
    trajectory_yaml_list = os.listdir(trajectory_folder)

    color = ['white', 'black', 'red', 'green', 'orange', 'purple','gray', 'brown', 'yellow', 'blue']
    used_color = color[0:len(trajectory_yaml_list)+2]
    cmap = ListedColormap(used_color)
    legend = []
    for index,baseline in enumerate(trajectory_yaml_list):
        legend.append(mpatches.Patch(color=used_color[index+2], label=baseline))
        baseline_path = trajectory_folder + '/' + baseline
        with open(baseline_path, 'r') as stream:
            yaml_data = yaml.load(stream, Loader=yaml.CLoader)

        if(quantity == None):
            for path in yaml_data["trajectories"]:
                for pose in path:
                    c = int((pose["position"][0] - info_yaml_data["origin_x"])/info_yaml_data["resolution"])
                    r = int((pose["position"][1] - info_yaml_data["origin_y"])/info_yaml_data["resolution"])

                    if c >= info_yaml_data["map_width"] or c < 0 : continue
                    if r >= info_yaml_data["map_height"] or r < 0 : continue
                    vis_trajectory_map[r][c] = index + 2
        else:
            for i in range(int(quantity)):
                path = yaml_data["trajectories"][i]
                for pose in path:
                    c = int((pose["position"][0] - info_yaml_data["origin_x"])/info_yaml_data["resolution"])
                    r = int((pose["position"][1] - info_yaml_data["origin_y"])/info_yaml_data["resolution"])

                    if c >= info_yaml_data["map_width"] or c < 0 : continue
                    if r >= info_yaml_data["map_height"] or r < 0 : continue
                    vis_trajectory_map[r][c] = index + 2

    fig, ax = plt.subplots()
    ax.imshow(vis_trajectory_map)
    ax.legend(handles=legend, prop={'size': 16})
    barrr = ax.imshow(vis_trajectory_map,cmap=cmap)

    figure_path = data_root + '/' + env + '/' + 'figure'
    if not os.path.exists(figure_path):
        os.makedirs(os.path.join(data_root + '/' + env, 'figure'))

    fig.savefig(figure_path + '/trajectory.png')

# Cell
def tortuosity(yaml_path):
    """
        Tortuosity calculated using the formula (arc length / chord length)
        Args:
            yaml_path (String) : The full path of yaml file.

        Returns:
            avg_tortuosity (Float) : The average tortuosity of all pathes in such yaml file.
    """
    dist_between_points = lambda x1, y1, x2, y2: ((x2-x1)**2 + (y2-y1)**2)**0.5

#     print(yaml_path)
    with open(yaml_path, 'r') as stream:
        yaml_data = yaml.load(stream, Loader=yaml.CLoader)

    total = 0.0
    for i_path, path in enumerate(yaml_data["trajectories"]):
        arc_len = 0.0
        for i_pose, pose in enumerate(path):
            if i_pose == 0: continue
            x1, y1 = path[i_pose-1]["position"][0], path[i_pose-1]["position"][1]
            x2, y2 = path[i_pose]["position"][0], path[i_pose]["position"][1]
            arc_len += dist_between_points(x1, y1, x2, y2)

        x_first, y_first = path[0]["position"][0], path[0]["position"][1]
        x_last, y_last = path[-1]["position"][0], path[-1]["position"][1]
        chord_len = dist_between_points(x_first, y_first, x_last, y_last)
        tor = arc_len / chord_len
#         print("The " ,i_path+1,"th path with tortuosity " ,tor)
        total += tor
    avg = total/len(yaml_data["trajectories"])
#     print("Avg tortuosity of all path: ", avg)
    return avg

# Cell
def map_size(data_root, env):
    """
        Real size of the environment.
        Args:
            data_root (String) : The path of folder env_exp_files. Default at 'Data/env_exp_files' .
            env (String) : Which environment was used.
        Returns:
            map_width (Int) : Map width (meter).
            map_height (Int) : Map width (meter).
    """
    yaml_path = data_root + '/' + env + '/' + 'environment_info' + '/' + 'grid_map_info.yaml'
    with open(yaml_path, 'r') as f:
        info_yaml_data = yaml.load(f, Loader=yaml.CLoader)
    map_width = info_yaml_data["map_width"]*info_yaml_data["resolution"]
    map_height = info_yaml_data["map_height"]*info_yaml_data["resolution"]
    print("Map Width: ",map_width)
    print("Map Height: ",map_height)
    return map_width, map_height

def avg_distance_between_pair(data_root, env, file_name=None):
    """
        Calculate average distance between start point and goal point.
        Args:
            data_root (String) : The path of folder env_exp_files. Default at 'Data/env_exp_files' .
            env (String) : Which environment was used.
            file_name (String) : The file stored all pairs.
        Returns:
            avg_dis (Float) : The average of the distance between Start and Goal points (meter).
    """
    dist_between_points = lambda x1, y1, x2, y2: ((x2-x1)**2 + (y2-y1)**2)**0.5

    if(file_name):
        pair_yaml_path = data_root + '/' + env + '/' + 'environment_info' + '/' + file_name +'.yaml'
    else:
        pair_yaml_path = data_root + '/' + env + '/' + 'environment_info' + '/' + env +'.yaml'
    with open(pair_yaml_path, 'r') as stream:
        yaml_data = yaml.load(stream, Loader=yaml.CLoader)

    total_dis = 0
    for pair in yaml_data["pairs"]:
        x1, y1 = pair["start"][0], pair["start"][1]
        x2, y2 = pair["goal"][0], pair["goal"][1]
        dis = dist_between_points(x1, y1, x2, y2)
        total_dis += dis

    avg_distance = total_dis/len(yaml_data["pairs"])
    print("The average of the distance between Start and Goal points: ", avg_distance)
    return avg_distance


# Cell
def random_color(number):
    '''
    generate random colors for each pair points and the number of color should not less than point's number.
    Args :
        number (int) : total number points.

    Returns :
        unique_color (list) : a list of different color.
    '''
    color=['#FFFFFF','#000000']
    if len(color)<number:
        for i in range (number):
            color.append("#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)]))
            i+=1
    unique_color = []
    for name in color:         # 1st loop
        if name not in unique_color:   # 2nd loop
            unique_color.append(name)
    if (len(unique_color)<number+2):
        unique_color = random_color(number)
    return unique_color

def draw_pair_points(data_root, env, file_name=None):
    '''
    Draw every pair points in the map from path : 'env'/environment_info/'env'.yaml .
    Args :
        data_root (String) : The path of folder env_exp_files. Default at 'Data/env_exp_files' .
        env (String) : Which environment was used.
        file_name (String) : The file stored all pairs.
    Returns:
        A figure (pairPoint.png) would be saved.
    '''
    # default path
    folder_map = 'npy' # The relative path to find the yaml file `bit.map`.
    map_file = 'obstacle_map.npy' # The file name of the 2D occupency grid map.
    info_file = 'grid_map_info.yaml' #The file name of the yaml file about 2D occupency grid map.

    # grid_map_info
    yaml_path = data_root + '/' + env + '/' + 'environment_info/' + info_file
    with open(yaml_path, 'r') as f:
        info_yaml_data = yaml.load(f, Loader=yaml.CLoader)
    # obstacle map
    map_path = data_root + '/' + '/' + env + '/npy/' + map_file
    pair_points_map = np.load(map_path)

    # start-goal pairs
    if(file_name):
        pair_yaml_path = data_root + '/' + env + '/' + 'environment_info' + '/' + file_name+'.yaml'
    else:
        pair_yaml_path = data_root + '/' + env + '/' + 'environment_info' + '/' + env+'.yaml'
    with open(pair_yaml_path, 'r') as stream:
        yaml_data = yaml.load(stream, Loader=yaml.CLoader)

    # color
    color=random_color(len(yaml_data["pairs"]))
    cmap=ListedColormap(color)
    #color_index change .npy map's color.
    color_index = 2

    for pair in yaml_data["pairs"]:
        x1, y1 = pair["start"][0], pair["start"][1]
        x2, y2 = pair["goal"][0], pair["goal"][1]

        c1 = int((x1 - info_yaml_data["origin_x"]) / info_yaml_data["resolution"])
        c2 = int((x2 - info_yaml_data["origin_x"]) / info_yaml_data["resolution"])
        r1 = int((y1 - info_yaml_data["origin_y"]) / info_yaml_data["resolution"])
        r2 = int((y2 - info_yaml_data["origin_y"]) / info_yaml_data["resolution"])

        # Line from interpolation
        delta_c = (c1-c2)*0.001
        delta_r = (r1-r2)*0.001
        for delta in range(1000):
            c = int( c2 + delta * delta_c )
            r = int( r2 + delta * delta_r )
            pair_points_map[r][c] = color_index
#         pair_points_map[r2][c2] = color_index

        color_index += 1

    fig, ax = plt.subplots()
    ax.set_title(f'{env} pair points map',color='r')
    ax.imshow(pair_points_map)
    barrr = ax.imshow(pair_points_map,cmap=cmap)

    figure_path = data_root + '/' + env + '/' + 'figure'
    if not os.path.exists(figure_path):
        os.makedirs(os.path.join( data_root + '/' + env, 'figure'))

    fig.savefig(figure_path + '/pairPoint.png')


# Cell
def plot_all_env(data_root, which_metric):
    '''
    Use sub-figure to visualize the metric of all environments together.
    You can plot 1.pairPoint 2.trajectory.
    you must run draw_pair_points() and plot_trajectory() to get all env's pairPoint.png and trajectory.png !
    Args:
        data_root (String) : Path. All directory under this path is an environment. Please check file structure in 03_diff_visualization.ipynb
        which_metric (String) : Metric name.
    Returns:
        None. A figure will pop up.

    '''
    envs = os.listdir(data_root)
    envs.remove('stage1')
    envs.remove('cave')
    img = []
    for i, env in enumerate(envs):
        img.append(Image.open(data_root+'/'+env+'/figure/'+which_metric+'.png'))
    fig, ((ax_1, ax_2, ax_3), (ax_4, ax_5, ax_6))= plt.subplots(nrows=2, ncols=3)
    fig.subplots_adjust(left=0,
                        bottom=0.1,
                        right=7,
                        top=5,
                        wspace=0,
                        hspace=0)
    ax_1.imshow(img[0])
    ax_1.axis("off")
    ax_2.imshow(img[1])
    ax_2.axis("off")
    ax_3.imshow(img[2])
    ax_3.axis("off")
    ax_4.imshow(img[3])
    ax_4.axis("off")
    ax_5.imshow(img[4])
    ax_5.axis("off")
    ax_6.imshow(img[5])
    ax_6.axis("off")

    plt.show()