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

__all__ = ['get_environments_metrics_traj', 'get_environments_metrics_pair']

# 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 numpy as np
from sklearn.datasets import load_iris
import matplotlib.pyplot as plt
from sklearn.manifold import MDS
from sklearn.preprocessing import MinMaxScaler
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
from mpl_toolkits.axes_grid1 import make_axes_locatable
from sklearn.preprocessing import normalize
plt.rcParams["figure.figsize"] = (16, 16)

# Cell
def get_environments_metrics_traj(data_root):
    """
    Calculate the average metric of all environment.
    Args:
        data_root (String) : Path. All directory under this path is an environment. Please check file structure in 03_diff_visualization.ipynb
    Returns:
        all_env (Dict) : A Dictionary like this
                            { 'garage': { 'avg_visibility': 15.987035875225928,
                                          'characteristic_dimension': 6.787662532508829,
                                          'closest_wall': 2.254822777080131,
                                          'dispersion': 1.2634011090573012},
                              'forest': { 'avg_visibility': 33.32919485818832,
                                          'characteristic_dimension': 26.731711305286364,
                                          'closest_wall': 3.6290804054099532,
                                          'dispersion': 0.6094594594594595}, ...}
    """
    envs = os.listdir(data_root)
    metrics = ['avg_visibility','characteristic_dimension','closest_wall','dispersion']

    all_env = dict()
    for env in envs:
        each_env = dict()
        # grid_map_info
        yaml_path = data_root + '/' + env + '/' + 'environment_info/' + 'grid_map_info.yaml'
        with open(yaml_path, 'r') as f:
            info_yaml_data = yaml.load(f, Loader=yaml.CLoader)

        # trajectory
        pair_yaml_path = data_root + '/' + env + '/' + 'trajectory' + '/' + env+'_TARE.yaml'
        with open(pair_yaml_path, 'r') as stream:
            pair_yaml_data = yaml.load(stream, Loader=yaml.CLoader)

        column_row = []
        for trajectory in pair_yaml_data["trajectories"]:
            for pose in trajectory:
                x = pose["position"][0]
                y = pose["position"][1]
                c = int((x - info_yaml_data["origin_x"]) / info_yaml_data["resolution"])
                r = int((y - info_yaml_data["origin_y"]) / info_yaml_data["resolution"])
                if [c, r] not in column_row: column_row.append( [c, r] )

        for metric in metrics:
            metric_data = np.load(data_root + '/' + env + '/' + 'npy' + '/' + metric+'.npy')
            norm_metric_data = metric_data

            total = 0.0
            for cr in column_row:
                if cr[1]>=norm_metric_data.shape[0] or cr[0]>=norm_metric_data.shape[1]: continue
                total += norm_metric_data[cr[1]][cr[0]]
            avg = total/len(column_row)

            each_env[metric] = avg
        each_env["tortuosity"] =  tortuosity('Data/env_exp_files/'+env+'/trajectory/'+env+'_TARE.yaml')
#         each_env["avg_distance"] = avg_distance_between_pair('Data/env_exp_files', env)
        all_env[env] = each_env
    return all_env

# Cell
def get_environments_metrics_pair(data_root):
    """
    Calculate the average metric of all environment.
    Args:
        data_root (String) : Path. All directory under this path is an environment. Please check file structure in 03_diff_visualization.ipynb
    Returns:
        all_env (Dict) : A Dictionary like this
                            { 'garage': { 'avg_visibility': 15.987035875225928,
                                          'characteristic_dimension': 6.787662532508829,
                                          'closest_wall': 2.254822777080131,
                                          'dispersion': 1.2634011090573012},
                              'forest': { 'avg_visibility': 33.32919485818832,
                                          'characteristic_dimension': 26.731711305286364,
                                          'closest_wall': 3.6290804054099532,
                                          'dispersion': 0.6094594594594595}, ...}
    """
    envs = os.listdir(data_root)
    metrics = ['avg_visibility','characteristic_dimension','closest_wall','dispersion']

    all_env = dict()
    for env in envs:
        each_env = dict()
        # grid_map_info
        yaml_path = data_root + '/' + env + '/' + 'environment_info/' + 'grid_map_info.yaml'
        with open(yaml_path, 'r') as f:
            info_yaml_data = yaml.load(f, Loader=yaml.CLoader)

        # start-goal pairs
        pair_yaml_path = data_root + '/' + env + '/' + 'environment_info' + '/' + env+'.yaml'
        with open(pair_yaml_path, 'r') as stream:
            pair_yaml_data = yaml.load(stream, Loader=yaml.CLoader)

        column_row = []
        for pair in pair_yaml_data["pairs"]:
            x1, y1 = pair["start"][0], pair["start"][1]
            x2, y2 = pair["goal"][0], pair["goal"][1]
            c1 = int((x1 - info_yaml_data["origin_x"]) / info_yaml_data["resolution"])
            c2 = int((x2 - info_yaml_data["origin_x"]) / info_yaml_data["resolution"])
            r1 = int((y1 - info_yaml_data["origin_y"]) / info_yaml_data["resolution"])
            r2 = int((y2 - info_yaml_data["origin_y"]) / info_yaml_data["resolution"])

            # Line from interpolation
            delta_c = (max(c1,c2) - min(c1,c2))/1000
            delta_r = (max(r1,r2) - min(r1,r2))/1000
            for delta in range(1000):
                c = int( min(c1,c2) + delta * delta_c )
                r = int( min(r1,r2) + delta * delta_r )
                if [c, r] not in column_row: column_row.append( [c, r] )

        for metric in metrics:
            metric_data = np.load(data_root + '/' + env + '/' + 'npy' + '/' + metric+'.npy')
            norm_metric_data = metric_data

            total = 0.0
            for cr in column_row:
                if cr[1]>=norm_metric_data.shape[0] or cr[0]>=norm_metric_data.shape[1]: continue
                total += norm_metric_data[cr[1]][cr[0]]
            avg = total/len(column_row)

            each_env[metric] = avg
        each_env["tortuosity"] =  tortuosity('Data/env_exp_files/'+env+'/trajectory/'+env+'_TARE.yaml')
#         each_env["avg_distance"] = avg_distance_between_pair('Data/env_exp_files', env)
        all_env[env] = each_env
    return all_env