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

__all__ = ['DifficultyMetrics', 'DifficultyVisualizer']

# Cell
class DifficultyMetrics:

  # map: C-space occupancy grid
  # path: list of points (row, col)
  # disp_radius: radius for dispersion
  def __init__(self, map, path, disp_radius):
    self.map = map
    self.rows = len(map)
    self.cols = len(map[0])
    self.axes = [(0, 1), (1, 1), (1, 0), (1, -1)] # vertical, horizontal, and 2 diagonals
    self.path = path
    self.radius = disp_radius

  # returns grid with the distance to closest obstacle at each point
  def closest_wall(self):
    dists = [[0 for i in range(self.cols)] for j in range(self.rows)]
    for r in range(self.rows):
      for c in range(self.cols):

        dists[r][c] = self._dist_closest_wall(r, c)

    return dists

  # returns grid with the average visibility at each point
  def avg_visibility(self):
    vis = [[0 for i in range(self.cols)] for j in range(self.rows)]
    for r in range(self.rows):
      for c in range(self.cols):
        vis[r][c] = self._avg_vis_cell(r, c)

    return vis

  # returns grid with the dispersion at each point
  # checks along 16 axes within the dispersion radius
  def dispersion(self):
    disp = [[0 for i in range(self.cols)] for j in range(self.rows)]
    for r in range(self.rows):
      for c in range(self.cols):
        disp[r][c] = self._cell_dispersion(r, c, self.radius)

    return disp

  # returns grid with the characteristic dimension at each point
  # characteristic dimension calculated in 2 directions for 4 axes
  def characteristic_dimension(self):
    cdr = [[0 for i in range(self.cols)] for j in range(self.rows)]
    for r in range(self.rows):
      for c in range(self.cols):
        if self.map[r][c] == 1:
          cdr[r][c] = 0

        cdr_min = self.rows + self.cols
        for axis in self.axes:
          cdr_min = min(cdr_min, self._distance(r, c, axis))

        cdr[r][c] = cdr_min

    return cdr

  # returns the distance along the axis in both directions, not including (r, c)
  def _distance(self, r, c, axis):
    if self.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 self.map[r_curr][c_curr] != 1:
        r_curr += move[0]
        c_curr += move[1]

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

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

    return dist


  # a and b are tuples (row, col)
  def _dist_between_points(self, a, b):
    return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)

  # path is list of points as tuples (row, col)
  # returns tortuosity calculated using the formula (arc length / chord length)
  def tortuosity(self):
    # calculate arc length
    arc_len = 0.0
    for i in range(1, len(self.path)):
      arc_len += self._dist_between_points(self.path[i - 1], self.path[i])

    # calculate chord length
    chord_len = self._dist_between_points(self.path[0], self.path[-1])
    return arc_len / chord_len

  # returns the cell dispersion within the radius
  def _cell_dispersion(self, r, c, radius):
    if self.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 >= self.rows or c_curr < 0 or c_curr >= self.cols:
          break

        if self.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

  # returns the average visibility at the point (r, c)
  def _avg_vis_cell(self, 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):
            break

          if self.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


  # bounds check
  def _in_map(self, r, c):
    return r >= 0 and r < self.rows and c >= 0 and c < self.cols

  # returns the distance to the closest obstacle at point (r, c)
  # returns 0 if self.map[r][c] is an obstacle, 1 if an adjacent non-diagonal cell is an obstacle, etc.
  def _dist_closest_wall(self, r, c):
    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 self.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) 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 (self.rows - 1) / 2

  # wrapper class for coordinates
  class Wrapper:

    def __init__(self, distance, row, col):
      self.dist = distance
      self.r = row
      self.c = col

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

  # returns a list of the unnormalized metrics, averaged over the path
  # metrics order: distance to closest wall, average visibility, dispersion,
  # characteristic dimension, and tortuosity
  def avg_all_metrics(self):
    result = []
    total = 0.0

    # closest wall
    total = 0.0
    for row, col in self.path:
      total += self._dist_closest_wall(row, col)
    result.append(total / len(self.path))

    # average visibility
    total = 0.0
    for row, col in self.path:
      total += self._avg_vis_cell(row, col)
    result.append(total / len(self.path))

    # dispersion
    total = 0.0
    for row, col in self.path:
      total += self._cell_dispersion(row, col, self.radius)
    result.append(total / len(self.path))

    # characteristic dimension
    total = 0.0
    char_dim_grid = self.characteristic_dimension()
    for row, col in self.path:
      total += char_dim_grid[row][col]
    result.append(total / len(self.path))

    # tortuosity
    tort = self.tortuosity()
    result.append(tort)

    return result

# Cell
class DifficultyVisualizer:
    def __init__(self, yaml_path, trajectory_yaml_list = []):
        self.yaml_path = yaml_path
        self.trajectory_yaml_list = trajectory_yaml_list
        with open(self.yaml_path, 'r') as stream:
            self.yaml_data = yaml.load(stream, Loader=yaml.CLoader)
        print(self.yaml_data)
        self.read_ply()


    def analyze_trajectory(self):
        for index, baseline in enumerate(self.trajectory_yaml_list):
            count = 0
            closest_wall = 0
            visibility = 0
            dispersion = 0
            characteristic_dimension = 0
            total_tortuosity = 0

            with open(baseline) as f: yaml_data = yaml.safe_load(f)
            for path in yaml_data:
                arc_len = 0.0
                for path_point_index, xy in enumerate(path):
                    c = int((xy[0] - self.origin_x)/self.resolution)
                    r = int((xy[1] - self.origin_y)/self.resolution)

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

                    closest_wall += self.total_closest_wall[r][c]
                    visibility += self.total_avg_visibility[r][c]
                    dispersion += self.total_dispersion[r][c]
                    characteristic_dimension += self.total_characteristic_dimension[r][c]
                    count += 1

                    # Tortuosity
                    if path_point_index==0: continue
                    arc_len += ((path[path_point_index][0] - path[path_point_index-1][0])**2 +
                                (path[path_point_index][1] - path[path_point_index-1][1])**2)**0.5
                chord_len = ((path[len(path)-1][0] - path[0][0])**2 +
                                (path[len(path)-1][1] - path[0][1])**2)**0.5
                tortuosity = arc_len / chord_len
                total_tortuosity += tortuosity


            print("------------------------- ",baseline, " -------------------------")
            print("avg closest_wall:", closest_wall/count)
            print("avg visibility:", visibility/count)
            print("avg dispersion:", dispersion/count)
            print("avg characteristic_dimension", characteristic_dimension/count)
            print("avg path tortuosity", total_tortuosity/len(yaml_data))


    def plot_trajectory(self):
        cmap = ListedColormap(['white', 'black', 'red', 'green'])
        vis_trajectory_map = self.map_bit.copy()
        for color_index,baseline in enumerate(self.trajectory_yaml_list):
            with open(baseline) as f: yaml_data = yaml.safe_load(f)

            for path in yaml_data:
                for xy in path:
                    c = int((xy[0] - self.origin_x)/self.resolution)
                    r = int((xy[1] - self.origin_y)/self.resolution)

                    if c >= self.map_height or c < 0 : continue
                    if r >= self.map_width or r < 0 : continue
                    vis_trajectory_map[r][c] = color_index + 2



        fig, ax = plt.subplots()
        ax.legend(self.trajectory_yaml_list)
        barrr = ax.imshow(vis_trajectory_map,cmap=cmap)
        fig.savefig('trajectory.png')


    def read_ply(self):
        # read ply file
        pc = o3d.io.read_point_cloud(self.yaml_data['source'])
        points_array = np.asarray((pc.points))

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

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

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

            self.map[r][c] += 1

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

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

    def plot_metrics(self):
        load_obstacle_map = self.map_bit  #np.load(self.yaml_data['source'][:-4]+'_bit.npy')
        shape = load_obstacle_map.shape
        diff = DifficultyMetrics(load_obstacle_map, [], disp_radius=1)
        self.total_obstacle_map = load_obstacle_map
        self.total_closest_wall = diff.closest_wall()
        self.total_avg_visibility = diff.avg_visibility()
        self.total_dispersion = diff.dispersion()
        self.total_characteristic_dimension = diff.characteristic_dimension()


        str_list = ["obstacle_map",\
                    "closest_wall",\
                    "avg_visibility",\
                    "dispersion",\
                    "characteristic_dimension"]
        cmap_list = ["binary","RdYlGn","RdYlGn","RdYlGn","Greys"]
        for i,item in enumerate([self.total_obstacle_map,self.total_closest_wall,self.total_avg_visibility, \
                                 self.total_dispersion,self.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(self.yaml_data['source'][:-4]+str_list[i]+".png")
            np.save(self.yaml_data['source'][:-4]+str_list[i]+".npy", np.asarray(item))