Analysis_side_SA.py

import os
import pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx
import matplotlib.patches as patches
import scipy.stats as sp


def save_result_data(Results=None):
    result_cols = ['Collision', 'braked', 'braked_pre', 'braked_post',
                   'accelerated', 'accelerated_pre', 'accelerated_post', 
                   'left_pre', 'left_post', 'right_pre', 'right_post', 'ego_pass_via_center', 'ego_pass_via_shoulder']
    # Drop old column
    if 'use_looming_reward' in Results.columns:
        Results = Results.drop(columns='use_looming_reward')

    # Add initial time headway
    try:
        Results['Initial THW'] = Results['x_tar_init'] / Results['v_tar_init']
    except:
        pass

    # Transfrom boolean columns to int
    bool_columns = Results.select_dtypes(include='bool').columns
    Results[bool_columns] = Results[bool_columns].astype(int)

    # Get useful columns 
    Ra = Results[(Results != 'Leeds').all(1)].to_numpy()
    No_EA = np.where(Ra == 'None')[0]
    Ra[Ra == 'None']     = 0.0
    Ra[Ra == 'Surprise'] = 1.0

    Ra = Ra.astype(float)

    # Remove columns with only a single value
    useful_columns = (np.nanstd(Ra, axis=0) > 0) | (np.in1d(Results.columns, result_cols) & (Ra.sum(0) > 0))

    Results = Results.iloc[:,useful_columns]

    # Remove columns which are linear combinations of other columns
    Results_copy = Results.copy()
    if 'EA_mode' in Results.columns:
        Results_copy.EA_mode = 1.0
        Results_copy.EA_mode.iloc[No_EA] = 0.0
    Results_copy = Results_copy.astype(float)
    if 'noise_pred_fac' in Results_copy.columns:
        Results_copy.noise_pred_fac = np.log(Results_copy.noise_pred_fac)
    linear_depend = Results_copy.corr().to_numpy() >= 0.9999

    # set main diagonal to False
    np.fill_diagonal(linear_depend, False)

    # Get connected subgraphs
    G = nx.Graph(linear_depend)
    connected_subgraphs = list(nx.connected_components(G))

    used_columns = []
    for subgraph in connected_subgraphs:
        used_columns.append(list(subgraph)[-1])

    # Extract unique columns
    useful_columns = np.zeros(Results_copy.shape[1], dtype=bool)
    useful_columns[used_columns] = True

    Results = Results.iloc[:,useful_columns]
    Results_copy = Results_copy.iloc[:,useful_columns]

    # Check which indices are actual results
    result_col = []
    menu_col = []
    for col in Results.columns:
        if col in result_cols:
            result_col.append(col)
        else:
            menu_col.append(col)

    Results = Results[menu_col + result_col]

    # Sort results by index
    Results = Results.sort_index()

    # Save the results
    Results.to_excel("Results_intersection_SA" + os.sep + 'Analysis_intersection_SA_leeds.xlsx', float_format="%.6f")
    return Results



def check_for_collisions(L, W, Center_tar, Center_ego, theta_tar, theta_ego):
     # Translate so ego is at origin
    Center_tar -= Center_ego

    # Rotate Center_tar so ego is aligned with x-axis
    Center_tar = np.stack((Center_tar[...,0] * np.cos(-theta_ego) - Center_tar[...,1] * np.sin(-theta_ego),
                           Center_tar[...,0] * np.sin(-theta_ego) + Center_tar[...,1] * np.cos(-theta_ego)), axis=-1)
    angle = theta_tar - theta_ego

    # Get the corner position of the target vehicle
    Corner = np.array([[-L/2, -W/2],
                       [ L/2, -W/2],
                       [ L/2,  W/2],
                       [-L/2,  W/2]])
    
    # Rotate the corners with theta_tar_adjust
    Corner = (Corner.T)[np.newaxis,np.newaxis]
    angle = angle[...,np.newaxis]

    Corner = np.stack((Corner[...,0,:] * np.cos(angle) - Corner[...,1,:] * np.sin(angle),
                       Corner[...,0,:] * np.sin(angle) + Corner[...,1,:] * np.cos(angle)), axis=-2)
    
    # Translate the corners to the center
    Corner += Center_tar[...,np.newaxis]

    # Check for each corner if any of its lines intersect with the ego vehicle (i.e., there is a collision)
    Line_above = (Corner[:,1:, 1] >  0.5 * W) & (Corner[:,:-1, 1] >  0.5 * W)
    Line_below = (Corner[:,1:, 1] < -0.5 * W) & (Corner[:,:-1, 1] < -0.5 * W)
    Line_left  = (Corner[:,1:, 0] < -0.5 * L) & (Corner[:,:-1, 0] < -0.5 * L)
    Line_right = (Corner[:,1:, 0] >  0.5 * L) & (Corner[:,:-1, 0] >  0.5 * L)

    # Definitely no colloision
    No_collision = (Line_above | Line_below | Line_left | Line_right)
    case_ind, interval_ind, corner_ind = np.where(~No_collision)

    Potential_line = np.stack((Corner[case_ind,interval_ind, :, corner_ind], 
                               Corner[case_ind,interval_ind + 1, :, corner_ind]), axis=-2)
    
    Px0 = Potential_line[:,0,0]
    Py0 = Potential_line[:,0,1]
    Pdx = Potential_line[:,1,0] - Potential_line[:,0,0]
    Pdy = Potential_line[:,1,1] - Potential_line[:,0,1]

    T_left  = (-L/2 - Px0) / Pdx
    T_right = ( L/2 - Px0) / Pdx
    T_above = ( W/2 - Py0) / Pdy
    T_below = (-W/2 - Py0) / Pdy

    T_hor_enter = np.minimum(T_left, T_right)
    T_hor_exit  = np.maximum(T_left, T_right)
    T_ver_enter = np.minimum(T_above, T_below)
    T_ver_exit  = np.maximum(T_above, T_below)

    T_enter = np.maximum(T_hor_enter, T_ver_enter)
    T_exit  = np.minimum(T_hor_exit, T_ver_exit)

    Collided = (T_enter < T_exit) & (T_enter < 1) & (T_exit >= 0)

    # Get collision times
    T_coll = np.ones((Corner.shape[0], Corner.shape[1] - 1, Corner.shape[3])) * np.nan
    T_coll[case_ind, interval_ind, corner_ind] = np.where(Collided, T_enter, np.nan) 

    # Find first collided corner
    T_coll = np.nanmin(T_coll, axis=-1)

    # Find the first collided timestep
    Collision_times = np.argmax(~np.isnan(T_coll), axis=1)
    T_coll = T_coll[np.arange(len(T_coll)), Collision_times] + Collision_times

    return T_coll

def get_stimuli_times(y_tar, lane_width, lf):
    # Get the time at which the target agent moved into the intersection
    y_radius = 10 + 0.5 * lane_width + lf - 0.01

    # Get the time that y_tar crosses -y_radius
    T_stimulus = np.ones(len(y_tar)) * np.nan
    for i_case in range(len(y_tar)):
        t_cand = np.where(y_tar[i_case,:] > -y_radius)[0]
        if len(t_cand) == 0:
            continue
        t_cand = t_cand[0]
        if t_cand == 0:
            T_stimulus[i_case] = 0.0
        else:
            y_0 = y_tar[i_case, t_cand - 1]
            y_1 = y_tar[i_case, t_cand]
            T_stimulus[i_case] = t_cand - 1 + ( -y_radius - y_0) / (y_1 - y_0)

    return T_stimulus


def get_steering_times(delta_ego, T_stimulus, L, dt=0.1):
    # Get response times for first brake, then for first left steer, and for first right steer
    
    # Get the curvature corresponding to the reference (steering wheel angle = 5 degrees)
    steering_wheel_angle_dec = 6
    curvature_steering_wheel_angle_ratio = 1 / 2717.75
    curvature_dec = steering_wheel_angle_dec * curvature_steering_wheel_angle_ratio

    # We know that for cuvature k = tan(delta) * cos(arctan(0.5* tan(delta))) / L ~ delta / L
    delta_dec = curvature_dec * L

    # Estimate first time that delta_ego > delta_dec, and theta_ego > 0
    T_left_cand = delta_ego > delta_dec
    # Only use candidates where the previous one is false
    T_left_cand[:,1:] &= ~T_left_cand[:,:-1]
    # For each candidate, find the first time that delta_ego > delta_dec, which is before T_left_cand.argmax()
    T_left = np.ones((len(delta_ego), max(1, T_left_cand.sum(-1).max())), dtype=int) * np.nan
    t_left_case, t_left_interval = np.where(T_left_cand)

    # Find the time that delta_ego == delta_dec before each t_left_interval
    j = -1
    i_case_old = 0
    for i_case, j_interval in zip(t_left_case, t_left_interval):
        if i_case_old == i_case:
            j += 1
        else:
            j = 0
        t_left_cand = np.where(delta_ego[i_case, :j_interval + 1] > delta_dec)[0][-1] - 1
        assert t_left_cand + 1 == j_interval, 'The last value before the candidate is not the candidate'
        delta_left_0 = delta_ego[i_case, t_left_cand]
        delta_left_1 = delta_ego[i_case, t_left_cand + 1]
        
        T_left[i_case, j] = t_left_cand + (delta_dec - delta_left_0) / (delta_left_1 - delta_left_0)
        
        i_case_old = i_case
    
    RT_left = (T_left - T_stimulus[:,np.newaxis]) * dt

    # Estimate first time that delta_ego < -delta_dec, and theta_ego < 0
    T_right_cand = delta_ego < -delta_dec
    # Only use candidates where the previous one is false
    T_right_cand[:,1:] &= ~T_right_cand[:,:-1]

    # For each candidate, find the first time that delta_ego < -delta_dec, which is before T_right_cand.argmax()
    T_right = np.ones((len(delta_ego), max(1,T_right_cand.sum(-1).max())), dtype=int) * np.nan
    t_right_case, t_right_interval = np.where(T_right_cand)

    # Find the time that delta_ego == -delta_dec before each t_right_interval
    j = -1
    i_case_old = 0
    for i_case, j_interval in zip(t_right_case, t_right_interval):
        if i_case_old == i_case:
            j += 1
        else:
            j = 0
        t_right_cand = np.where(delta_ego[i_case, :j_interval + 1] < -delta_dec)[0][-1] - 1
        assert t_right_cand + 1 == j_interval, 'The last value before the candidate is not the candidate'
        delta_right_0 = delta_ego[i_case, t_right_cand]
        delta_right_1 = delta_ego[i_case, t_right_cand + 1]
        
        # Check if delta or theta made this a fullfilled criteria
        T_right[i_case, j] = t_right_cand + (-delta_dec - delta_right_0) / (delta_right_1 - delta_right_0)
        
        i_case_old = i_case

    RT_right = (T_right - T_stimulus[:,np.newaxis]) * dt

    if (RT_right == np.inf).any():
        print('There are infinite values in RT_right')
    return RT_left, RT_right


def get_brake_and_acc_times(a_ego, T_stimulus, dt = 0.1):
    acc_threshold = 1 
    # Estimate first time that a_ego < -acc_threshold, and theta_ego < 0
    T_brake_cand = a_ego < -acc_threshold
    # Only use candidates where the previous one is false
    T_brake_cand[:,1:] &= ~T_brake_cand[:,:-1]

    # For each candidate, find the first time that a_ego < -acc_threshold, which is before T_brake_cand.argmax()
    T_brake = np.ones((len(a_ego), max(1,T_brake_cand.sum(-1).max())), dtype=int) * np.nan
    t_brake_case, t_brake_interval = np.where(T_brake_cand)

    # Find the time that a_ego == -acc_threshold before each t_brake_interval
    j = -1
    i_case_old = 0

    for i_case, j_interval in zip(t_brake_case, t_brake_interval):
        if i_case_old == i_case:
            j += 1
        else:
            j = 0
        t_brake_cand = np.where(a_ego[i_case, :j_interval + 1] < -acc_threshold)[0][-1] - 1
        assert t_brake_cand + 1 == j_interval, 'The last value before the candidate is not the candidate'
        # Test if we use actual accelerations (constant timewise instead of linear interpolation)
        a_brake_0 = a_ego[i_case, t_brake_cand]
        a_brake_1 = a_ego[i_case, t_brake_cand + 1]
        
        t_brake = t_brake_cand + (acc_threshold - a_brake_0) / (a_brake_1 - a_brake_0)
        T_brake[i_case, j] = (t_brake + 4 * (t_brake_cand + 1)) / 5 
        
        i_case_old = i_case
    
    RT_brake = (T_brake - T_stimulus[:,np.newaxis]) * dt

    # Estimate first time that a_ego > acc_threshold, and theta_ego > 0
    T_acc_cand = a_ego > acc_threshold
    # Only use candidates where the previous one is false
    T_acc_cand[:,1:] &= ~T_acc_cand[:,:-1]

    # For each candidate, find the first time that a_ego > acc_threshold, which is before T_acc_cand.argmax()
    T_acc = np.ones((len(a_ego), max(1,T_acc_cand.sum(-1).max())), dtype=int) * np.nan
    t_acc_case, t_acc_interval = np.where(T_acc_cand)

    # Find the time that a_ego == acc_threshold before each t_acc_interval
    j = -1
    i_case_old = 0

    for i_case, j_interval in zip(t_acc_case, t_acc_interval):
        if i_case_old == i_case:
            j += 1
        else:
            j = 0
        t_acc_cand = np.where(a_ego[i_case, :j_interval + 1] > acc_threshold)[0][-1] - 1
        assert t_acc_cand + 1 == j_interval, 'The last value before the candidate is not the candidate'
        a_acc_0 = a_ego[i_case, t_acc_cand]
        a_acc_1 = a_ego[i_case, t_acc_cand + 1]
        
        t_acc = t_acc_cand + (acc_threshold - a_acc_0) / (a_acc_1 - a_acc_0)
        T_acc[i_case, j] = 0.5 * (t_acc + t_acc_cand + 1)
        
        i_case_old = i_case
    
    RT_acc = (T_acc - T_stimulus[:,np.newaxis]) * dt

    return RT_brake, RT_acc








#####################################################################################################################

# Load Meta data
Meta_file = "Results_intersection_SA" + os.sep + "Setups_intersection_SA.xlsx"
Meta = pd.read_excel(Meta_file, index_col=0, keep_default_na=False)

# Load simple meta data
Meta_simple_file = "Results_intersection_SA" + os.sep + "Setups_simple_intersection_SA.xlsx"
Meta_simple = pd.read_excel(Meta_simple_file, index_col=0, keep_default_na=False)

# Go throught the results directory
result_files = os.listdir("Results_intersection_SA")

# Prepare results
Results = Meta_simple.copy()

# Transform looming rewards into actual values
LR_values = {'V2': (1, False),
             'V3': (0.5, False),
             'V4': (0.25, False),
             'V5': (1, True),
             'V6': (0.5, True),
             'V7': (0.25, True)}

# Disentangle reward function
Results['LR_std_fac'] = np.nan
Results['LR_collision'] = False
if 'use_looming_reward' in Results.columns:
    for exp_index in Results.index:
        meta = Results.loc[exp_index]
        LR = meta['use_looming_reward']
        if LR in LR_values:
            Results.loc[exp_index, 'LR_std_fac'] = LR_values[LR][0]
            Results.loc[exp_index, 'LR_collision'] = LR_values[LR][1]

# Add columns for the results
Results['Collision'] = np.nan
Results['braked'] = np.nan
Results['steered'] = np.nan


all_RT_first = []
Experiment = []
Collision = []
Steered = []
Braked = []

for number, result_file in enumerate(result_files):
    # Exclude .xlsx files
    if ".xlsx" in result_file:
        continue

    # Exclude pdf files
    if ".pdf" in result_file:
        continue

    # Exclude svg files
    if ".svg" in result_file:
        continue

    # Exclude npy files
    if ".npy" in result_file:
        continue

    # Exclude odt files
    if ".odt" in result_file:
        continue

    # Exclude odt files
    if ".png" in result_file:
        continue

    if ".csv" in result_file:
        continue

    print('Analyzing: ' + result_file + ' (' + str(number + 1) + ')')

    # Get experiment number
    exp_index = int(result_file.split('_')[-1])
    meta = Meta.loc[exp_index]
    Experiment.append(exp_index)

    # Load the results
    folder_path = "Results_intersection_SA" + os.sep + result_file + os.sep
    with open(folder_path + result_file + '.pkl', 'rb') as f:
        data = pickle.load(f)

    # Get driven trajectories
    Eta = data['eta']
    Traj_ego = Eta[...,[0,1,2,3,4]] # [x, y, theta, delta, v]
    Traj_tar = Eta[...,[5,6,7,8,9]] # [x, y, theta, delta, v]

    # Split Traj and Ego target
    x_ego, y_ego, theta_ego, delta_ego, v_ego = Traj_ego[...,0], Traj_ego[...,1], Traj_ego[...,2], Traj_ego[...,3], Traj_ego[...,4]
    x_tar, y_tar, theta_tar, delta_tar, v_tar = Traj_tar[...,0], Traj_tar[...,1], Traj_tar[...,2], Traj_tar[...,3], Traj_tar[...,4]
    a_tar, w_tar = Eta[...,10], Eta[...,11]
    a_ego, w_ego = data['a_cont'][0,...,0].T, data['a_cont'][0,...,1].T

    # get scenario parameters
    dt = meta['dt']
    lane_width = meta['lane_width']
    vehicle_width = meta['d']
    vehicle_length = (meta['lf'] + meta['lr'])

    # Get stimulus indices
    T_stimulus = get_stimuli_times(y_tar, lane_width, meta['lf'])

    # CHeck for collisions between ego and target agent (more precise check due to angles, with interpolation)
    Center_tar = np.stack((x_tar, y_tar), axis=-1)
    Center_ego = np.stack((x_ego, y_ego), axis=-1)

    # get collisions times
    T_coll = check_for_collisions(vehicle_length, vehicle_width, Center_tar, Center_ego, theta_tar, theta_ego)
    t_coll = (T_coll - T_stimulus) * dt
    # Get collisions
    collision = np.isfinite(T_coll)

    left_lane = ((y_tar < - 0.5 * (lane_width - vehicle_width)) | (y_tar > 0.5 * (3 * lane_width - vehicle_width))).any(-1)
    # collision |= left_lane

    # Check if agent steered/braked
    steered = np.abs(y_ego).max(-1) > 0.5 * lane_width
    braked = a_ego.min(-1) <= -1.0

    Collision.append(collision)
    Steered.append(steered)
    Braked.append(braked)

    # Get steering times
    RT_left, RT_right = get_steering_times(delta_ego, T_stimulus, vehicle_length, dt)

    # Get brake and acceleration times
    RT_brake, RT_acc = get_brake_and_acc_times(a_ego, T_stimulus, dt)

    # This is in order: Brake, Acc, Left, Right
    RT_firsts = np.stack((RT_brake[:,0], RT_acc[:,0], RT_left[:,0], RT_right[:,0]), axis=-1)

    # Only copnsider pre collision responses
    t_coll[np.isnan(t_coll)] = np.inf

    past_coll_response = RT_firsts >= t_coll[:,np.newaxis]
    RT_firsts[past_coll_response] = np.nan

    # Save the results
    all_RT_first.append(RT_firsts)

    # Save the results
    Results.loc[exp_index, 'Collision'] = collision.mean()
    Results.loc[exp_index, 'braked'] = braked.mean()
    Results.loc[exp_index, 'steered'] = steered.mean()


Results = save_result_data(Results)



# Leeds unique settings
Experiment = np.array(Experiment)
all_RT_first = np.array(all_RT_first)
Collision = np.stack(Collision, 0)
Steered = np.stack(Steered, 0)
Braked = np.stack(Braked, 0)

experiment_sorted = np.argsort(Experiment) 

all_RT_first = all_RT_first[experiment_sorted]
Collision = Collision[experiment_sorted]
Steered = Steered[experiment_sorted]
Braked = Braked[experiment_sorted]


# Load the actual GT data
GT_data = pd.read_csv("Intersection_data.csv")
GT_data['braked'] = np.in1d(GT_data['Response'], ['BS', 'B'])
GT_data['steered'] = np.in1d(GT_data['Response'], ['BS', 'S'])
# Remove Response column
GT_data = GT_data.drop(columns=['Response'])

# Calculate first response times for all models
# From Meta simple, get columns that have more than one unique value
use_columns = []
for col in Meta_simple.columns:
    if Meta_simple[col].nunique() > 1:
        use_columns.append(col)
Settings = Meta_simple[use_columns]
# Replace 'v_tar' column with 'Scenario' column
Settings['Scenario'] = 'RS'
Settings.loc[Meta['v_tar'] > 1.0, 'Scenario'] = 'RNS'
if 'v_ego' in Settings.columns:
    Settings = Settings.drop(columns=['v_ego'])
if 'v_ego_des' in Settings.columns:
    Settings = Settings.drop(columns=['v_ego_des'])
if 'x_ego' in Settings.columns:
    Settings = Settings.drop(columns=['x_ego'])

if 'v_tar' in Settings.columns: 
    Settings = Settings.drop(columns=['v_tar'])
if 'y_tar' in Settings.columns:
    # Overwrite with the different values for reasonable_stop_acceleration
    y_radius = 10 + 0.5 * Meta['lane_width'] + Meta['lf']
    help = - y_radius - Meta['y_tar']
    t_added = help / Meta['v_tar']
    help2 = t_added - Meta['ttc_trigger']
    d_stop = help2 * Meta['v_tar']
    a_assumed = (Meta['v_tar'] ** 2) / (2 * d_stop)

    # Round a_assumed to 3 decimals
    a_assumed = np.round(a_assumed, 3)
    Settings['a_assumed'] = a_assumed
    Settings = Settings.drop(columns=['y_tar'])

# Copy the columns with RS to have the different a_assumed scenarios
Settings['Experiment'] = np.arange(len(Settings))
unique_a_assumed = Settings[Settings['Scenario'] == 'RNS']['a_assumed'].unique()
Settings.loc[(Settings['Scenario'] == 'RS'), 'a_assumed'] = unique_a_assumed[0]
Settings_RS = Settings[(Settings['Scenario'] == 'RS')].copy()
for a_assumed in unique_a_assumed[1:]:
    Settings_RS_copy = Settings_RS.copy()
    Settings_RS_copy['a_assumed'] = a_assumed
    Settings = pd.concat([Settings, Settings_RS_copy], axis=0)
Settings = Settings.reset_index(drop=True)

# Move experiemnt doubling to other data
experiment_id = Settings['Experiment'].to_numpy()
all_RT_first = all_RT_first[experiment_id]
Collision = Collision[experiment_id]
Steered = Steered[experiment_id]
Braked = Braked[experiment_id]
Settings = Settings.drop(columns=['Experiment'])


# Get unique settings excluding rel_target
unique_settings, unique_indices = np.unique(Settings.drop(columns=['Scenario']).to_numpy().astype(str), axis = 0, return_inverse=True)

Results_settings = pd.DataFrame(unique_settings, columns=Settings.drop(columns=['Scenario']).columns)

delta_t = 0.2
# Get the unique rel_target values in meta_sample
n_cases = 32
RT_first = np.ones((len(unique_settings), 2, 3 * n_cases, 2)) * np.nan
Behavior_model = np.zeros((len(unique_settings), 2, 3 * n_cases, 5))

RT_data = np.zeros((2, 26, 2))
Behavior_data = np.zeros((2, 26, 5))
for i in range(len(unique_settings)):
    ind = np.where(unique_indices == i)[0]
    settings = Settings.iloc[ind]

    if i == 0:
        for j, scenario in enumerate(['RS', 'RNS']):
            gt_data = GT_data[GT_data['Scenario'] == scenario]
            num_samples = len(gt_data)

            Behavior_data[j,:num_samples,0] = gt_data['Collision'].to_numpy().astype(float)

            brake_bool = gt_data['braked'].to_numpy().astype(bool)
            steer_bool = gt_data['steered'].to_numpy().astype(bool)
            Behavior_data[j,:num_samples,1] = (brake_bool & (~steer_bool)).astype(float) # Brake only
            Behavior_data[j,:num_samples,2] = ((~brake_bool) & steer_bool).astype(float) # Steer only
            Behavior_data[j,:num_samples,3] = (brake_bool & steer_bool).astype(float)  # Both
            Behavior_data[j,:num_samples,4] = (~brake_bool & ~steer_bool).astype(float) # No response



            Behavior_data[j,num_samples:] = np.nan

            rt_data = gt_data['RT'].to_numpy().astype(float)
            rt_data_complex = np.full((len(rt_data), 2), np.nan)
            rt_data_complex[brake_bool,0] = rt_data[brake_bool]
            rt_data_complex[steer_bool,1] = rt_data[steer_bool]
            RT_data[j,:num_samples] = rt_data_complex
            RT_data[j,num_samples:] = np.nan


    for j, scenario in enumerate(['RS', 'RNS']):
        # Get comparable indices
        ind_rel = np.where(settings['Scenario'] == scenario)[0]
        i_rel = np.array(settings.index[ind_rel])
        num_samples = len(i_rel) * n_cases

        # Extract values
        all_rt = all_RT_first[i_rel].reshape(-1, 4)
        RT_first[i,j,:num_samples,0] = all_rt[:,0]
        RT_first[i,j,:num_samples,1] = np.nanmin(all_rt[:,2:], axis=-1)
        rt_scenario = RT_first[i,j, :num_samples]
        rt_min = np.nanmin(rt_scenario, axis=-1, keepdims=True)
        first_reaction = rt_scenario <= rt_min + delta_t
        rt_scenario[~first_reaction] = np.nan
        RT_first[i,j,:num_samples] = rt_scenario

        # Extract behavior types
        Behavior_model[i,j,:num_samples,0] = Collision[i_rel].flatten()
        braked_bool = Braked[i_rel].flatten().astype(bool)
        steered_bool = Steered[i_rel].flatten().astype(bool)
        Behavior_model[i,j,:num_samples,1] = (braked_bool & (~steered_bool)).astype(float) # Brake only
        Behavior_model[i,j,:num_samples,2] = ((~braked_bool) & steered_bool).astype(float) # Steer only
        Behavior_model[i,j,:num_samples,3] = (braked_bool & steered_bool).astype(float)  # Both
        Behavior_model[i,j,:num_samples,4] = (~braked_bool & ~steered_bool).astype(float) # No response
        Behavior_model[i,j,num_samples:] = np.nan


# Bootstrap the JSD values
n_bootstraps = 10000
JSD_response = np.zeros((len(unique_settings), 2, n_bootstraps))
JSD_collision = np.zeros((len(unique_settings), 2, n_bootstraps))
for i in range(len(unique_settings)):
    for j in range(2):
        available_data = ~np.isnan(Behavior_data[j,:,0])
        available_model = ~np.isnan(Behavior_model[i,j,:,0])
        behavior_data = Behavior_data[j,available_data] # shape (n_samples, 5)
        behavior_model = Behavior_model[i,j,available_model] # shape (64, 5)

        data_coll = behavior_data[:,0]
        model_coll = behavior_model[:,0]

        data_response = behavior_data[:,1:].argmax(-1)
        model_response = behavior_model[:,1:].argmax(-1)
        for b in range(n_bootstraps):
            data_b = np.random.choice(data_response, size=len(data_response), replace=True).astype(int)
            model_b = np.random.choice(model_response, size=len(model_response), replace=True).astype(int)

            data_c = np.random.choice(data_coll, size=len(data_coll), replace=True).astype(int)
            model_c = np.random.choice(model_coll, size=len(model_coll), replace=True).astype(int)

            # Get probabilities
            prob_data_c = np.bincount(data_c, minlength=2) / len(data_c)
            prob_model_c = np.bincount(model_c, minlength=2) / len(model_c)
            prob_data_b = np.bincount(data_b, minlength=4) / len(data_b)
            prob_model_b = np.bincount(model_b, minlength=4) / len(model_b)

            # Get Jensen shannon divergence
            M_c = 0.5 * (prob_data_c + prob_model_c)
            M_b = 0.5 * (prob_data_b + prob_model_b)

            Data_c_nonzero = prob_data_c > 0
            Model_c_nonzero = prob_model_c > 0
            Data_b_nonzero = prob_data_b > 0
            Model_b_nonzero = prob_model_b > 0

            KLD_c_data = np.zeros(M_c.shape)
            KLD_c_model = np.zeros(M_c.shape)
            KLD_b_data = np.zeros(M_b.shape)
            KLD_b_model = np.zeros(M_b.shape)

            KLD_c_data[Data_c_nonzero] = prob_data_c[Data_c_nonzero] * np.log(prob_data_c[Data_c_nonzero] / M_c[Data_c_nonzero])
            KLD_c_model[Model_c_nonzero] = prob_model_c[Model_c_nonzero] * np.log(prob_model_c[Model_c_nonzero] / M_c[Model_c_nonzero])
            KLD_b_data[Data_b_nonzero] = prob_data_b[Data_b_nonzero] * np.log(prob_data_b[Data_b_nonzero] / M_b[Data_b_nonzero])
            KLD_b_model[Model_b_nonzero] = prob_model_b[Model_b_nonzero] * np.log(prob_model_b[Model_b_nonzero] / M_b[Model_b_nonzero])

            KLD_c_data = KLD_c_data.sum(-1)
            KLD_c_model = KLD_c_model.sum(-1)
            KLD_b_data = KLD_b_data.sum(-1)
            KLD_b_model = KLD_b_model.sum(-1)

            JSD_collision[i,j,b] = 0.5 * (KLD_c_data + KLD_c_model)
            JSD_response[i,j,b] = 0.5 * (KLD_b_data + KLD_b_model)

# average over scenarios
jsd_response = JSD_response.mean(-2)
jsd_collision = JSD_collision.mean(-2)

# Get the Kullback-Leibler Divergence
## Get RT KS test (assume n = 20, and alpha = 0.05)

# estimate the maximum difference between linearly interpolated CDF
# Potential: First response, so everything larger than 0.2s after first response is nan
RT_model = np.sort(RT_first, axis = -2) # shape (n, 2, 64, 2)
RT_data = np.sort(RT_data, axis = -2)   # shape (2, 26, 2)

num_model_samples = np.isfinite(RT_model).sum(-2, keepdims=True)
num_data_samples = np.isfinite(RT_data).sum(-2, keepdims=True)

num_bootstraps = 10000
Wasserstein_D = np.zeros((len(unique_settings), 2, 2, num_bootstraps))

for i in range(len(unique_settings)):
    for j in range(2):
        for k in range(2):
            rt_model = RT_model[i,j,:,k]
            rt_data = RT_data[j,:,k]

            available_model = np.isfinite(rt_model)
            available_data = np.isfinite(rt_data)

            rt_model = rt_model[available_model]
            rt_data = rt_data[available_data]

            if any(available_model) and any(available_data):

                for b in range(num_bootstraps):
                    rt_data_b = np.random.choice(rt_data, size=len(rt_data), replace=True)
                    rt_model_b = np.random.choice(rt_model, size=len(rt_model), replace=True)
                    # Get wasserstein distance
                    Wasserstein_D[i,j,k,b] = sp.wasserstein_distance(rt_model_b, rt_data_b)
            else:
                Wasserstein_D[i,j,k,:] = np.nan


# Get the mean values over scenarios
wasserstein_d = Wasserstein_D.mean((1,2))

# Get 5, 25, 50, 75, 95 percentiles over bootstraps
jsd_quantiles = np.percentile(jsd_collision, [5, 25, 50, 75, 95], axis=-1).T
wasserstein_d_quantiles = np.percentile(wasserstein_d, [5, 25, 50, 75, 95], axis=-1).T 

# Write JSD and WD to csv file
quatiles = np.concatenate((jsd_quantiles, wasserstein_d_quantiles), axis=1)
columns = [
    'JSDI_5', 'JSDI_25', 'JSDI_50', 'JSDI_75', 'JSDI_95',
    'WDI_5', 'WDI_25', 'WDI_50', 'WDI_75', 'WDI_95'
]
results_csv = pd.DataFrame(quatiles, columns=columns)
# In unique settings, find for each column the value seen most often
norm_settings = []
cols = []
for col in Meta_simple.columns:
    if col not in Results_settings.columns:
        continue
    cols.append(col)
    values = Results_settings[col].values
    unique, count = np.unique(values, return_counts=True)
    most_frequent = unique[np.argmax(count)]
    norm_settings.append(most_frequent)
norm_settings = np.array(norm_settings)

# Find norm setting index
norm_index = np.where((np.abs((Results_settings[cols].to_numpy().astype(float) - norm_settings.astype(float))) < 1e-5).all(-1))[0]

# Go through each column again:
for j, col in enumerate(cols):
    values = Results_settings[col].values
    if len(np.unique(values)) == 1:
        continue

    # Find indices in column where this value is not the norm value
    diff_indices = np.where(values != norm_settings[j])[0]
    if len(norm_index) > 0:
        used_indices = np.array([norm_index[0]] + diff_indices.tolist())
    else:
        used_indices = diff_indices
    try:
        used_values = values[used_indices].astype(float)

        i_argsort = np.argsort(used_values)
        used_indices = used_indices[i_argsort]
        used_values = used_values[i_argsort]
    except:
        pass

    results_col = results_csv.iloc[used_indices]
    results_col['param'] = used_values

    results_col.to_csv("Results_intersection_SA" + os.sep + 'Analysis_settings_' + col + '.csv')

import Analysis_SA