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

__all__ = ['plot_obstacle_map', 'plot_trajectory', 'load_csv', 'load_world']

# Cell
import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt
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
from queue import Queue, PriorityQueue
import random
plt.rcParams["figure.figsize"] = (16, 16)

# Cell
def plot_obstacle_map(data_root, env, folder, filename):
    '''
        Plot obstacle map.
        Args:
            data_root (String) : The path of folder env_exp_files. Default at 'Data/env_exp_files' .
            env (String) : Which environment was used.
            folder (String) : The relative path to find the map file.
            filename (String) : The file name of numpy array map.

        Returns:
            obstacle_map (array) : obstacle map would also be saved in .npy files and visualized in .png files.

    '''
    map_path = data_root + '/' + env + '/' + folder + '/' + filename
    load_obstacle_map = np.load(map_path)
    obstacle_map = load_obstacle_map

    fig, ax = plt.subplots()
    barrr = ax.imshow(obstacle_map,cmap="binary", interpolation='nearest')
    ax.set_title("obstacle_map")
    fig.colorbar(barrr, ax=ax)
    npy_file_name = data_root + '/' + env + '/npy/' + "obstacle_map"+".npy"
    png_name = data_root + '/' + env + '/figure/' + "obstacle_map"+".png"
    fig.savefig(png_name)
    np.save(npy_file_name, np.asarray(load_obstacle_map))
    return obstacle_map

# Cell
def plot_trajectory(data_root, env, folder_map, map_file, folder_path, folder_info, info_file):
    """
        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.
            folder_map (String) : The relative path to find the yaml file `bit.map`.
            map_file (String) : The file name of the 2D occupency grid map.
            folder_path (String) : The relative path to find the yaml file of trajectories.
            folder_info (String) : The relative path to find the yaml file of information about 2D occupency grid map.
            info_file (String) : The file name of the yaml file about 2D occupency grid map.

        Returns:
            A figure (trajectory.png) would be saved.
    """

    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)

    cmap = ListedColormap(['white', 'black', 'red', 'green'])
    legend = []
    for index,baseline in enumerate(trajectory_yaml_list):
        legend.append(mpatches.Patch(color=['white', 'black', 'red', 'green'][index+2], label=baseline))

        baseline_path = trajectory_folder + '/' + baseline
        with open(baseline_path, 'r') as stream:
            yaml_data = yaml.load(stream, Loader=yaml.CLoader)

        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

    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 load_csv(env, csv_file):
    '''
    load csv file into numpy array.
    Args :
        env (String) : Which environment was used.
        csv_files (String) : csv file path.

    Returns :
        array : pair points in numpy array.
    '''
    path = 'Data/env_exp_files/'+ env + '/environment_info/'+ csv_file
    with open(path) as file_name:
        array = np.loadtxt(file_name, delimiter=",")
    return array

# Cell
from os import listdir
from os.path import isfile, join

# Cell
def load_world(path):
    have_found=[]
    with open(path,'r') as f:
        for line in f.readlines():
            if ("model://" in line) :
                if not line in have_found:
                    print(line)
                    have_found.append(line)
        f.close()