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

__all__ = ['get_environments_metrics_traj', 'get_metric_values', 'Normalization_1D', 'visualize_1D', 'visualize_2D']

# 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')
        all_env[env] = each_env
    return all_env

# Cell
def get_metric_values(all_env):
    """
    Calculate MAX, MIN, mean & std. for all metrics.
    Args
        all_env (Dict) : A Dictionary like this
                            { 'garage': { 'avg_visibility': 15.987035875225928,
                                          'characteristic_dimension': 6.787662532508829,
                                          'closest_wall': 2.254822777080131,
                                          'dispersion': 1.2634011090573012,
                                          'tortuosity': 1.291879835480693},
                              'forest': { 'avg_visibility': 33.32919485818832,
                                          'characteristic_dimension': 26.731711305286364,
                                          'closest_wall': 3.6290804054099532,
                                          'dispersion': 0.6094594594594595,
                                          'tortuosity': 1.291879835480693}, ...}
    Returns:
        None. But printtn the results.
    """
    metrics = ['avg_visibility', 'characteristic_dimension', 'closest_wall', 'dispersion', 'tortuosity']
    all_data = [[] for i in range(len(metrics))]
    for env_key in all_env:
        for i,metric in enumerate(metrics):
            all_data[i].append( np.array(all_env[env_key][metric]) )
    for i in range(len(all_data)):
        print(metrics[i])
        print('min=', round(np.min(all_data[i]),2),end=' ')
        print('max=', round(np.max(all_data[i]),2),end=' ')
        print('mean=', round(np.mean(all_data[i]),2),end=' ')
        print('std=', round(np.std(all_data[i]),2),end=' \n')

# Cell
def Normalization_1D(arr):
    """
    Normalization. Make the MAX in the array to be 1 and the MIN to be 0.
    Args
        arr (array) : 1D array.
    Returns:
        norm_arr (array) : 1D array.
    """
    MAX = np.max(arr)
    MIN = np.min(arr)

    norm_arr = []
    for i in range(len(arr)):
        norm_i = float( (arr[i]-MIN) / (MAX - MIN) )
        norm_arr.append(norm_i)
    return np.array( norm_arr )

def visualize_1D(all_env):
    """
    Visualize 1D distribution of metrics. All value been normalized to 0~1.
    Args
        all_env (Dict) : A Dictionary like this
                            { 'garage': { 'avg_visibility': 15.987035875225928,
                                          'characteristic_dimension': 6.787662532508829,
                                          'closest_wall': 2.254822777080131,
                                          'dispersion': 1.2634011090573012,
                                          'tortuosity': 1.291879835480693},
                              'forest': { 'avg_visibility': 33.32919485818832,
                                          'characteristic_dimension': 26.731711305286364,
                                          'closest_wall': 3.6290804054099532,
                                          'dispersion': 0.6094594594594595,
                                          'tortuosity': 1.291879835480693}, ...}
    Returns:
        None. But plot the figure.
    """
    metrics = ['avg_visibility', 'characteristic_dimension', 'closest_wall', 'dispersion', 'tortuosity']
    colors = ['red','green','yellow','grey','purple','orange','pink','cyan','silver']
    env_names = []
    all_data = [[] for i in range(len(metrics))]
    for env_key in all_env:
        env_names.append(env_key)
        for i,metric in enumerate(metrics):
            all_data[i].append( np.array(all_env[env_key][metric]) )
    all_data = [all_data[i] for i in range(len(all_data))]

    norm_all_data = [ Normalization_1D(all_data[i]) for i in range(len(all_data)) ]
    plt.figure()
    for i in range(len(norm_all_data)):
        plt.subplot(str(510+i))
        plt.title(metrics[i])
        for j,env_name in enumerate(env_names):
            plt.scatter([norm_all_data[i][j]],[0],color=colors[j],s=300 )

    if not os.path.exists('Data'):
        os.makedirs(os.path.join('Data'))
    plt.savefig("Data/diff_metrics_visualize_1D.png")
    print(env_names)
    print(colors)


# Cell
def visualize_2D(all_env):
    """
    Visualize 2D distribution of metrics.
    Consider all combinations of two out of five metrics. Therefore, here are 10 figures.
    All value been normalized to 0~1.
    Args
        all_env (Dict) : A Dictionary like this
                            { 'garage': { 'avg_visibility': 15.987035875225928,
                                          'characteristic_dimension': 6.787662532508829,
                                          'closest_wall': 2.254822777080131,
                                          'dispersion': 1.2634011090573012,
                                          'tortuosity': 1.291879835480693},
                              'forest': { 'avg_visibility': 33.32919485818832,
                                          'characteristic_dimension': 26.731711305286364,
                                          'closest_wall': 3.6290804054099532,
                                          'dispersion': 0.6094594594594595,
                                          'tortuosity': 1.291879835480693}, ...}
    Returns:
        None. But plot the figures.
    """
    metrics = ['avg_visibility', 'characteristic_dimension', 'closest_wall', 'dispersion', 'tortuosity']
    env_names = []
    colors = []
    all_data = [[] for i in range(len(metrics))]
    for env_key in all_env:
        env_names.append(env_key)
        if env_key=='cave' or env_key=='BARN_3x3':colors.append('green')
        else:colors.append('red')

        for i,metric in enumerate(metrics):
            all_data[i].append( np.array(all_env[env_key][metric]) )
    all_data = [all_data[i] for i in range(len(all_data))]
    norm_all_data = [ Normalization_1D(all_data[i]) for i in range(len(all_data)) ]

    if not os.path.exists('Data'): os.makedirs(os.path.join('Data'))
    plt.figure()
    for i in range(5):
        for j in range(i+1,5):
            fig, ax = plt.subplots()
            ax.scatter(norm_all_data[i],norm_all_data[j],s=300,color=colors)
            for k, font in enumerate (env_names):
                plt.annotate(font, (norm_all_data[i][k],norm_all_data[j][k]),size=20)
            ax.set_xlabel(metrics[i])
            ax.set_ylabel(metrics[j])
            ax.grid()
            fig.savefig('Data/'+metrics[i]+'_'+metrics[j]+'.png')
