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

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

# 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
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)

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

    def crop_sub_map(self,r,c,full_occu_map):
        """
        Crop a sub-map with fixed size (31*31) from the full occupancy grid map.
        The pixel full_occu_map[r][c] will be the center of sub-map, that is crop_map[15][15].
        Padding with 1, which means consider as obstacle if a pixel we need in sub-map in out of the full occupancy grid map.

        Args:
            full_occu_map (2D array) : The full occupancy grid map.
            r(int) : index in row in the full occupancy grid map. .
            c(int) : index in column in the full occupancy grid map. .
        Returns:
            crop_map (2D array) : An 31*31 sub-map from full_occu_map.
        """
        crop_map = np.ones((31,31))
        for sub_r in range(-15,16):
            for sub_c in range(-15,16):
                if (r+sub_r < 0) or (r+sub_r >=full_occu_map.shape[0]): continue
                if (c+sub_c < 0) or (c+sub_c >=full_occu_map.shape[1]): continue
                crop_map[sub_r+15][sub_c+15] = full_occu_map[r+sub_r][c+sub_c]
        return crop_map

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

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

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

    def characteristic_dimension(self,full_occu_map):
        """
            Calculate the average dispersion of an occupancy grid map.
            Characteristic dimension calculated in 2 directions for 4 axes
            Args:
                full_occu_map (2D array) : The full occupancy grid map.
            Returns:
                The average of characteristic dimension at each point.
            """

        cdr = [[0 for i in range(len(full_occu_map[0]))] for j in range(len(full_occu_map))]
        for r in range(len(full_occu_map)):
            for c in range(len(full_occu_map[0])):
                if full_occu_map[r][c] == 1:
                    cdr[r][c] = 0
                    continue
                crop_map = self.crop_sub_map(r, c, full_occu_map)
                cdr_min = ((15*2 + 15*2) ** 0.5) * 2 # diagonal
                for axis in [(0, 1), (1, 1), (1, 0), (1, -1)]:
                    cdr_min = min(cdr_min, self._distance(axis, crop_map))
                cdr[r][c] = cdr_min
        return cdr

    def _distance(self, axis, crop_map):
        """
        Args:
            crop_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)
        """
        r = 15
        c = 15
        if crop_map[r][c] == 1: return -1# is obstalce

        # 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 crop_map[r_curr][c_curr] != 1:
                r_curr += move[0]
                c_curr += move[1]

                # if not in map
                if (r_curr+move[0]>=31 or c_curr+move[1]>=31 or r_curr+move[0]<0 or c_curr+move[1]<0 ): break
                # add the distance traveled to the total
                if crop_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, crop_map):
        """
            Args:
                crop_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.
            """
        radius=1
        r = 15
        c = 15
        if crop_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(crop_map) or c_curr < 0 or c_curr >= len(crop_map[0]):break
                if crop_map[r_curr][c_curr] == 1:wall = True
                if move[0] == 2 or move[1] == 2:
                    count += 2 # count the in-between axes as two steps
                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_pixel(self, crop_map):
        """
            Args:
                crop_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).
            """
        r = 15
        c = 15
        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 ( r_curr>=31 or c_curr>=31 or r_curr<0 or c_curr<0 ) : break # not in map
                    if crop_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, full_occu_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(full_occu_map) and c >= 0 and c < len(full_occu_map[0])


    def _dist_closest_wall(self, crop_map):
        """
            Args:
                crop_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 crop_map[r][c] is an obstacle, 1 if an adjacent non-diagonal cell is an obstacle, etc.
            """
        pq = PriorityQueue()
        r = 15
        c = 15
        first_wrapper = self.Wrapper(0, r, c)
        pq.put(first_wrapper)
        visited = {(r, c) : first_wrapper}
        while not pq.empty():
            point = pq.get()
            if crop_map[point.r][point.c] == 1: return point.dist # found an obstacle!
            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 (row<31 and row>=0 and col<31 and col>=0) 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(crop_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/' + 'obstacle_map.npy'
    np.save(npy_file_name, map_bit)
    plt.imshow(map_bit)
    plt.colorbar()
    plt.show()
    return map_bit


# Cell

def calculate_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:
            metric_list (array) : The result of metrics which order in [`obstacle map`,`closest wall`,`visibility`, `dispersion`,`characteristic dimension`].

    """
    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"]

    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'))

    for i,item in enumerate([total_obstacle_map,total_closest_wall,total_avg_visibility, \
                             total_dispersion,total_characteristic_dimension]):
        npy_file_name = data_root + '/' + env + '/npy/' + str_list[i]+".npy"
        np.save(npy_file_name, np.asarray(item))
    metric_list = [total_obstacle_map, total_closest_wall, total_avg_visibility, total_dispersion, total_characteristic_dimension]
    return metric_list

def plot_metric(data_root, env, metric_list):
    """
        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) : Name of the environment.
            metric_list (array) : The result of metrics which order in [`obstacle map`,`closest wall`,`visibility`, `dispersion`,`characteristic dimension`].

        Returns:
            None : Metrics visualization be saved in .png files. Full path is    data_root + '/' + env + '/figure/' + "XXX.png"
    """

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

    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(metric_list):
        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)

        png_name = data_root + '/' + env + '/figure/' + str_list[i]+".png"
        fig.savefig(png_name)

