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

__all__ = ['get_environments_metrics', 'metrics_multidimensional_scaling', 'plot_MDS_result']

# 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

plt.rcParams["figure.figsize"] = (16, 16)

# Cell
def get_environments_metrics(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

        all_env[env] = each_env
    return all_env

def metrics_multidimensional_scaling(all_env):
    """
    Use MDS to immerse the dissimilarities.
    Args
        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}, ...}
    Returns:
        env_data_2d (array) : Results of MDS
        env_text_name (List) : A List of environment name String.
        env_target (List) : A List of labels about if such environment could train from scratch.
    """
    env_data = []
    env_target = []
    env_text_name = []
    for env_key in all_env:
        env_text_name.append(env_key)
        data = []
        for env_metric in all_env[env_key]:
            data.append(all_env[env_key][env_metric])
        env_data.append(data)
        if env_key=="cave" or env_key=="stage1": env_target.append(1)
        else: env_target.append(0)

    scaler = MinMaxScaler()
    env_data_scaled = scaler.fit_transform(env_data)
    mds = MDS(n_components=2,random_state=0)
    env_data_2d = mds.fit_transform(env_data)
    return env_data_2d, env_text_name, env_target

def plot_MDS_result(env_data_2d, env_text_name, env_target):
    """
    Plot 2D scatter of the results of MDS.
    Args:
        env_data_2d (array) : Results of MDS
        env_text_name (List) : A List of environment name String.
        env_target (List) : A List of labels about if such environment could train from scratch.
    Returns:
        None. A figure will pop up.
    """
    colors = ['red', 'green']
    target_names = ['can\'t training_from_scratch', 'can training_from_scratch']
    for i in np.unique(env_target):
        subset = env_data_2d[env_target == i]
        x = [row[0] for row in subset]
        y = [row[1] for row in subset]
        plt.scatter(x,y,c=colors[i],label=target_names[i],s=250)
    for i, font in enumerate (env_text_name):
        plt.annotate(font, (env_data_2d[i][0],env_data_2d[i][1]),size=20)

    plt.grid()
    plt.title('Metrics Multidimensional Scaling',size=22)
    plt.tick_params(axis='x', labelsize=16)
    plt.tick_params(axis='y', labelsize=16)
    plt.legend()
    plt.show()