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

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

# Cell
import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
import yaml
import math
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(yaml_path):
    """
        Generate the 2D occupency grid map by 3D pointcloud map.
        Args:
            yaml_path (String) : The file path of the yaml file for map information.
        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
    print(yaml_path)
    with open(yaml_path, 'r') as stream:
        yaml_data = yaml.load(stream, Loader=yaml.CLoader)
        print(yaml_data)
    pc = o3d.io.read_point_cloud(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

    np.save(yaml_data['source'][:-4]+'_bit.npy', map_bit)
    plt.imshow(map_bit)
    plt.colorbar()
    plt.show()
    return map_bit


# Cell

def plot_metric(map_path):
    """
        Calculate metrics for the input map.
        Args:
            map_path (String) : The file path of the .npy file for the occupency grid map.
        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.
    """
    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"]
    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)
        fig.savefig(str_list[i]+".png")
        np.save(str_list[i]+".npy", np.asarray(item))

    return total_obstacle_map, total_closest_wall, total_avg_visibility, total_dispersion, total_characteristic_dimension