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

__all__ = ['DifficultyMetrics', 'read_ply', 'plot_metric', 'plot_trajectory', 'avg_distance']

# 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
plt.rcParams["figure.figsize"] = (16, 16)

# Cell
class DifficultyMetrics:
      def __init__(self):
        pass

      def closest_wall(self, input_map):
        """
        Calculate the avg Distance to Closest Obstacle of an occupancy grid map.
        Args:
            input_map (2D array) : occupancy grid map.
        Returns:
            Avg Distance to Closest Obstacle of all cell in map.
        """
        dists = [[0 for i in range(len(input_map[0]))] for j in range(len(input_map))]
        for r in range(len(input_map)):
          for c in range(len(input_map[0])):
            dists[r][c] = self._dist_closest_wall(r, c, input_map)

        return dists


      def avg_visibility(self, input_map):
        """
        Calculate the average visibility of an occupancy grid map.
        Args:
            input_map (2D array) : occupancy grid map.
        Returns:
            The average visibility at each point.
        """
        vis = [[0 for i in range(len(input_map[0]))] for j in range(len(input_map))]
        for r in range(len(input_map)):
          for c in range(len(input_map[0])):
            vis[r][c] = self._avg_vis_cell(r, c, input_map)

        return vis


      def dispersion(self, input_map):
        """
        Calculate the average dispersion of an occupancy grid map.
        Checks along 16 axes within the dispersion radius.
        Args:
            input_map (2D array) : occupancy grid map.
        Returns:
            The average dispersion at each point.
        """
        disp = [[0 for i in range(len(input_map[0]))] for j in range(len(input_map))]
        for r in range(len(input_map)):
          for c in range(len(input_map[0])):
            disp[r][c] = self._cell_dispersion(r, c, input_map)

        return disp


      def characteristic_dimension(self,input_map):
        """
        Calculate the average dispersion of an occupancy grid map.
        Characteristic dimension calculated in 2 directions for 4 axes
        Args:
            input_map (2D array) : occupancy grid map.
        Returns:
            The average of characteristic dimension at each point.
        """
        cdr = [[0 for i in range(len(input_map[0]))] for j in range(len(input_map))]
        for r in range(len(input_map)):
          for c in range(len(input_map[0])):
            if input_map[r][c] == 1:
              cdr[r][c] = 0

            cdr_min = len(input_map) + len(input_map[0])
            for axis in [(0, 1), (1, 1), (1, 0), (1, -1)]:
              cdr_min = min(cdr_min, self._distance(r, c, axis, input_map))

            cdr[r][c] = cdr_min

        return cdr


      def _distance(self, r, c, axis, input_map):
        """
        Args:
            input_map (2D array) : occupancy grid map.
            r(int) : index in row.
            c(int) : index in column.
            axis(int) : a list such as [0,1], indicate the directions.
        Returns:
            The distance along the axis in both directions, not including (r, c)
        """
        if input_map[r][c] == 1:
          return -1

        # check axis in both directions
        reverse_axis = (axis[0] * -1, axis[1] * -1)
        dist = 0
        for move in [axis, reverse_axis]:
          r_curr = r
          c_curr = c

          # move in axis direction until an obstacle is found
          while input_map[r_curr][c_curr] != 1:
            r_curr += move[0]
            c_curr += move[1]

            if not self._in_map(r_curr, c_curr, input_map):
              break

            # add the distance traveled to the total
            if input_map[r_curr][c_curr] != 1:
              dist += math.sqrt(move[0] ** 2 + move[1] ** 2)

        return dist



      def _dist_between_points(self, a, b):
        """
        Args:
            a and b are tuples : (row, col)
        Returns:
            The distance between points
        """
        return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)



      def _cell_dispersion(self, r, c, input_map, radius=1):
        """
        Args:
            input_map (2D array) : occupancy grid map.
            r(int) : index in row.
            c(int) : index in column.
            radius(int) : cell size, default as 1.
        Returns:
            The cell dispersion within the radius.
        """
        if input_map[r][c] == 1:
          return -1

        axes_wall = []
        # four cardinal, four diagonal, and one in between each (slope +- 1/2 or +-2)
        for move in [(0, 1), (1, 2), (1, 1), (2, 1), (1, 0), (2, -1), (1, -1), (1, -2), (0, -1), (-2, -1), (-1, -1), (-1, -2), (-1, 0), (-2, 1), (-1, 1), (-1, 2)]:
          count = 0
          wall = False
          r_curr = r
          c_curr = c
          while count < radius and not wall:
            r_curr += move[0]
            c_curr += move[1]

            if r_curr < 0 or r_curr >= len(input_map) or c_curr < 0 or c_curr >= len(input_map[0]):
              break

            if input_map[r_curr][c_curr] == 1:
              wall = True

            # count the in-between axes as two steps
            if move[0] == 2 or move[1] == 2:
              count += 2
            else:
              count += 1

          if wall:
            axes_wall.append(True)
          else:
            axes_wall.append(False)

        # count the number of changes in this cell's field of view
        change_count = 0
        for i in range(len(axes_wall)-1):
          if axes_wall[i] != axes_wall[i+1]:
            change_count += 1

        if axes_wall[0] != axes_wall[15]:
          change_count += 1

        return change_count


      def _avg_vis_cell(self, r, c, input_map):
        """
        Args:
            input_map (2D array) : occupancy grid map.
            r(int) : index in row.
            c(int) : index in column.
        Returns:
            The average visibility at the point (r, c).
        """
        total_vis = 0.0
        num_axes = 0
        for r_move in [-1, 0, 1]:
          for c_move in [-1, 0, 1]:
            if r_move == 0 and c_move == 0:
              continue

            this_vis = 0
            r_curr = r
            c_curr = c
            wall_found = False
            while not wall_found:
              if not self._in_map(r_curr, c_curr, input_map):
                break

              if input_map[r_curr][c_curr] == 1:
                wall_found = True
              else:
                this_vis += math.sqrt((r_move ** 2) + (c_move ** 2))

              r_curr += r_move
              c_curr += c_move

            total_vis += this_vis
            num_axes += 1

        return total_vis / num_axes



      def _in_map(self, r, c, input_map):
        """
        Bounds check.
        Args:
            r(int) : index in row.
            c(int) : index in column.
        Returns:
            A Boolean means that if the cell is in bounded.
        """
        return r >= 0 and r < len(input_map) and c >= 0 and c < len(input_map[0])


      def _dist_closest_wall(self, r, c, input_map):
        """
        Args:
            input_map (2D array) : occupancy grid map.
            r(int) : index in row.
            c(int) : index in column.
        Returns:
            The distance to the closest obstacle at point (r, c).
            Returns 0 if input_map[r][c] is an obstacle, 1 if an adjacent non-diagonal cell is an obstacle, etc.
        """
        pq = PriorityQueue()
        first_wrapper = self.Wrapper(0, r, c)
        pq.put(first_wrapper)
        visited = {(r, c) : first_wrapper}

        while not pq.empty():
          point = pq.get()
          if input_map[point.r][point.c] == 1: # found an obstacle!
            return point.dist
          else:
            # enqueue all neighbors if they are in the map and have not been visited
            for row in range(point.r - 1, point.r + 2):
              for col in range(point.c - 1, point.c + 2):
                if self._in_map(row, col, input_map) and (row, col) not in visited:
                  dist = math.sqrt((row - r) ** 2 + (col - c) ** 2)
                  neighbor = self.Wrapper(dist, row, col)
                  pq.put(neighbor)
                  visited[(row, col)] = neighbor

        # in case the queue is empty before a wall is found (shouldn't happen),
        # the farthest a cell can be from a wall is half the board, since the top and bottom rows are all walls
        return (len(input_map) - 1) / 2


      class Wrapper:
        """ wrapper class for coordinates """
        def __init__(self, distance, row, col):
          self.dist = distance
          self.r = row
          self.c = col

        def __lt__(self, value):
          return self.dist < value.dist


# Cell

def read_ply(data_root, env, folder, filename):
    """
        Generate the 2D occupency grid map by 3D pointcloud 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 yaml file.
            filename (String) : The file name of yaml.

            Finally, the full path of the target file would be "data_root + '/' + env + '/' + folder + '/' + filename".
        Returns:
            map_bit (array) : The occupency grid map in the form of an numpy array. It would also be saved in .npy file and visualized in .png file.
    """
    # read ply file
    yaml_path = data_root + '/' + env + '/' + folder + '/' + filename
    print(yaml_path)
    with open(yaml_path, 'r') as stream:
        yaml_data = yaml.load(stream, Loader=yaml.CLoader)
    pc = o3d.io.read_point_cloud(data_root + '/' + env + '/' + folder + '/' + yaml_data['source'])

    points_array = np.asarray((pc.points))

    map_height, map_width = yaml_data['map_height'], yaml_data['map_width']
    map = np.zeros((map_height, map_width))
    origin_x, origin_y = yaml_data['origin_x'], yaml_data['origin_y']
    resolution = yaml_data['resolution']
    for i in range(len(points_array)):
        if points_array[i][2] < yaml_data['z_max'] : continue
        if points_array[i][2] > yaml_data['z_min'] : continue

        c = int((points_array[i][0] - origin_x)/resolution)
        r = int((points_array[i][1] - origin_y)/resolution)

        if c >= map_width or c < 0 :
            continue
        if r >= map_height or r < 0 :
            continue

        map[r][c] += 1

    map_bit = np.zeros(map.shape)
    for c in range(map_width):
        for r in range(map_height):
            map_bit[r][c] = 1 if map[r][c] > yaml_data['threshold'] else 0

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

    npy_file_name = data_root + '/' + env + '/npy/' + 'bit.npy'
    np.save(npy_file_name, map_bit)
    plt.imshow(map_bit)
    plt.colorbar()
    plt.show()
    return map_bit


# Cell

def plot_metric(data_root, env, folder, filename):
    """
        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 (String) : The relative path to find the yaml file.
            filename (String) : The file name of yaml.

            Finally, the full path of the .npy file for the occupency grid map would be "data_root + '/' + env + '/' + folder + '/' + filename".

        Returns:
            5 map_bit (array) : The result of metrics which order in [`obstacle map`,`closest wall`,`visibility`, `dispersion`,`characteristic dimension`].
                                They 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)
    shape = load_obstacle_map.shape
    diff = DifficultyMetrics()
    total_obstacle_map = load_obstacle_map
    total_closest_wall = diff.closest_wall(load_obstacle_map)
    total_avg_visibility = diff.avg_visibility(load_obstacle_map)
    total_dispersion = diff.dispersion(load_obstacle_map)
    total_characteristic_dimension = diff.characteristic_dimension(load_obstacle_map)

    str_list = ["obstacle_map",\
                "closest_wall",\
                "avg_visibility",\
                "dispersion",\
                "characteristic_dimension"]
    cmap_list = ["binary","RdYlGn","RdYlGn","RdYlGn","Greys"]

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


    for i,item in enumerate([total_obstacle_map,total_closest_wall,total_avg_visibility, \
                             total_dispersion,total_characteristic_dimension]):
        fig, ax = plt.subplots()
        barrr = ax.imshow(item,cmap=cmap_list[i], interpolation='nearest')
        ax.set_title(str_list[i])
        fig.colorbar(barrr, ax=ax)
        npy_file_name = data_root + '/' + env + '/npy/' + str_list[i]+".npy"
        png_name = data_root + '/' + env + '/figure/' + str_list[i]+".png"
        fig.savefig(png_name)
        np.save(npy_file_name, np.asarray(item))

    return total_obstacle_map, total_closest_wall, total_avg_visibility, total_dispersion, total_characteristic_dimension

# 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 avg_distance(data_root, env, folder_info, csv_file, 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_info (String) : The relative path to find the yaml file of information about 2D occupency grid map.
            csv_file (String) : The file name of the csv file about Start and Goal points.

        Returns:
            avg_dis (Float) : The average of the distance between Start and Goal points.
            map_width (Int) : Map width (meter).
            map_height (Int) : Map width (meter).
    """

    csv_path = data_root + '/' + env + '/' + folder_info + '/' + csv_file
    csv_data = np.genfromtxt(csv_path, delimiter=',')

    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)

    total_dis = 0
    for i in range(len(csv_data)):
        if i%2==1: continue
        x1, y1 = csv_data[i][0], csv_data[i][1]
        x2, y2 = csv_data[i+1][0], csv_data[i+1][1]

        dis = ((x2-x1)**2 + (y2-y1)**2) ** 0.5
        total_dis += dis

    avg = 2*total_dis/len(csv_data)

    print("----------",env,"----------")
    print("Start and Goal points pairs: ", len(csv_data)/2)
    print("The average of the distance between Start and Goal points: ",avg)
    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 avg, map_width, map_height
