# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/lib_nbs/00_utils.ipynb.

# %% auto 0
__all__ = ['isBetween_c_Vec', 'coord_mod', 'isBetween_c_Vec_nAgents', 'get_encounters', 'get_config', 'get_policy',
           'get_performance', 'get_opt']

# %% ../nbs/lib_nbs/00_utils.ipynb 3
import numpy as np

# %% ../nbs/lib_nbs/00_utils.ipynb 5
def isBetween_c_Vec(a, b, c, r):
    """
    Checks whether point c is crossing the line formed with point a and b.

    Parameters
    ----------
    a : tensor, shape = (1,2)
        Previous position.
    b : tensor, shape = (1,2)
        Current position.
    c : tensor, shape = (Nt,2)
        Positions of all targets.
    r : int/float
        Target radius.
    
    Returns
    -------
    mask : array of boolean values
        True at the indices of found targets.

    """
    
    mask = np.array([True]*c.shape[0])
    
    #exclude the targets whose vertical projection of the vector c-a w.r.t. the vector b-a is larger than the target radius.
    mask[np.argwhere(np.abs(np.cross(b-a, c-a))/np.linalg.norm(b-a) > r)] = False    
    
    #exclude the targets whose scalar product is negative (they are on the other side of the step direction)
    dotproduct = (c[:, 0] - a[0]) * (b[0] - a[0]) + (c[:, 1] - a[1])*(b[1] - a[1])
    mask[np.argwhere(dotproduct < 0)] = False
    
    #exclude the targets that are beyond the step.
    squaredlengthba = (b[0] - a[0])*(b[0] - a[0]) + (b[1] - a[1])*(b[1] - a[1])
    mask[np.argwhere(dotproduct > squaredlengthba)] = False
    
    return mask

# %% ../nbs/lib_nbs/00_utils.ipynb 6
def coord_mod(coord1, coord2, mod):
    """
    Computes the distance difference between two coordinates, in a world with size 'mod'
    and periodic boundary conditions.

    Parameters
    ----------
    coord1 : value, np.array, tensor (can be shape=(n,1))
        First coordinate.
    coord2 : np.array, tensor -- shape=(1,1)
        Second coordinate, substracted from coord1.
    mod : int
        World size.

    Returns
    -------
    diff_min : float
        Distance difference (with correct sign, not absolute value).

    """
    diff = np.remainder(coord1 - coord2, mod)
    diff_min = np.minimum(diff, mod-diff)
    
    diff_min[diff_min != diff] = -diff_min[diff_min != diff]
    
    return diff_min


# %% ../nbs/lib_nbs/00_utils.ipynb 7
def isBetween_c_Vec_nAgents(a, b, c, r):
    """
    Checks whether point c is crossing the line formed with point a and b. Code to run several agents in parallel.

    Parameters
    ----------
    a : array, shape = (n,2)
        Previous position of all n agents.
    b : array, shape = (n,2)
        Current position of all n agents.
    c : array, shape = (Nt,2)
        Positions of all targets.
    r : float
        Target radius
    
    Returns
    -------
    mask : array of boolean values, shape = (Nt, n)
        True at the indices of found targets.

    """
    
    mask = np.ones([c.shape[0], a.shape[0]],dtype=bool)
    
    #exclude the targets whose vertical projection of the vector c-a w.r.t. the vector b-a is larger than the target radius.
    for i in range(a.shape[0]):
        mask[np.argwhere(np.abs(np.cross(b[i]-a[i], c-a[i]))/np.linalg.norm(b[i]-a[i]) > r),i] = False    
    
    #exclude the targets whose scalar product is negative (they are on the other side of the step direction)
    dotproduct = (np.transpose(np.tile(c[:,0],(a.shape[0],1)))- a[:,0]) * (b[:,0] - a[:,0]) + (np.transpose(np.tile(c[:,1],(a.shape[0],1)))- a[:,1])*(b[:,1] - a[:,1])
    for coord in np.argwhere(dotproduct < 0):
        mask[coord[0]][coord[1]] = False
    
    #exclude the targets that are beyond the step.
    squaredlengthba = (b[:,0] - a[:,0])*(b[:,0] - a[:,0]) + (b[:,1] - a[:,1])*(b[:,1] - a[:,1])
    for coord in np.argwhere(dotproduct > squaredlengthba):
        mask[coord[0]][coord[1]] = False
    
    return mask

# %% ../nbs/lib_nbs/00_utils.ipynb 8
def get_encounters(agent_previous_pos, agent_pos, target_positions, L, r):
    """
    Considering the agent walks, it checks whether the agent finds a target while walking the current step.
    Code to run several agents in parallel.
    
    Parameters
    ----------
    agent_previous_pos : array, shape = (n,2)
        Position of the n agents before taking the step.
    agent_pos : array, shape = (n,2)
        Position of the n agents.
    target_positions : array, shape = (Nt,2)
        Positions of the targets.
    L : int
        World size.
    r : float
        Radius of the targets.
    
    Returns
    -------
    mask_encounters : array of boolean values, shape = (Nt, n)
        True at the indices of found targets.

    """
    
    mask_encounters = isBetween_c_Vec_nAgents(agent_previous_pos, agent_pos, target_positions, r)
    distance = np.sqrt((agent_previous_pos[:,0]-np.transpose(np.tile(target_positions[:,0],(agent_previous_pos.shape[0],1))))**2 +\
                       (agent_previous_pos[:,1]-np.transpose(np.tile(target_positions[:,1],(agent_previous_pos.shape[0],1))))**2)          
    
    #if there is more than one encounter, pick the one nearest the agent's position.
    for ag, num_encounters in enumerate(np.sum(mask_encounters,axis=0)):
        if num_encounters > 1:
            min_distance_masked = np.argmin(distance[:, ag][mask_encounters[:, ag]])
            first_encounter = np.arange(len(target_positions))[mask_encounters[:, ag]][min_distance_masked]
            one_encounter_mask = np.zeros(len(target_positions), dtype=bool)
            one_encounter_mask[first_encounter] = True
            
            mask_encounters[:, ag] = one_encounter_mask
    
    if np.sum(mask_encounters) >  agent_previous_pos.shape[0]:
        print('One agent is getting more than one target')
    
    return mask_encounters

# %% ../nbs/lib_nbs/00_utils.ipynb 10
import yaml
import numpy as np
import pandas as pd

# %% ../nbs/lib_nbs/00_utils.ipynb 11
def get_config(config, config_path = 'configurations/learning/'):
    """
    
    Get the configuration file of the given experiment and config name (e.g. exp_0).
    
    Parameters
    ----------
    config : str
        Config name (e.g. exp_0)
    experiment : str
        Experiment name (e.g. 'learning')
    config_path : str
        path to configurations

    Returns
    -------
    config_dict : dict
        Dictionary with the parameters of the loaded configuration.
    """
    
    path = config_path + config 
    with open(path, 'r') as f:
        config_dict = yaml.safe_load(f)
    return config_dict

# %% ../nbs/lib_nbs/00_utils.ipynb 12
def get_policy(results_path, agent, episode):
    """
    Gets the policy of an agent at a given episode.

    Parameters
    ----------
    results_path : str
        Path of the folder from which to extract the data.
    agent : int
        Agent index.
    episode : int
        Episode.

    Returns
    -------
    y : list
        Policy.

    """
    memories = np.load(results_path + 'memory_agent_'+str(agent)+'_episode_'+str(episode)+'.npy')
    num_states = len(memories[0,0,:])
    
    y = []
    for state_index in range(num_states):
        y.append(memories[ 0, 0, state_index] / np.sum(memories[0, :,state_index]))
       
    return y


# %% ../nbs/lib_nbs/00_utils.ipynb 13
def get_performance(results_path, agent_list, episode_list):
    """
    Extract data with the efficiencies obtained in the postlearning analysis.

    Parameters
    ----------
    results_path : str
        Path of the folder from which to extract the data.
    agent_list : list
        List with the agent indices.
    episode_list : list
        List with the episodes.

    Returns
    -------
    post_training_performance : np.array, shape=(len(agent_list), len(episode_list))
        Average performances obtained by the agents in the postlearning analysis.
    sem : np.array, shape=(len(agent_list), len(episode_list))
        Standard deviation of the mean.

    """
    post_training_performance = np.zeros([len(agent_list), len(episode_list)])
    sem = np.zeros([len(agent_list), len(episode_list)])
    
    for ind, ag in enumerate(agent_list):
        for e_index, episode in enumerate(episode_list):
            perf = np.load(results_path + 'performance_post_training_agent_'+str(ag)+'_episode_'+str(episode)+'.npy')
            post_training_performance[ind, e_index] = np.mean(perf)
            sem[ind, e_index] = np.std(perf) / np.sqrt(len(perf))
            
    return post_training_performance, sem

# %% ../nbs/lib_nbs/00_utils.ipynb 14
def get_opt(path, df):
    """
    Get the highest efficiency obtained by the benchmark models and the corresponding parameters.

    Parameters
    ----------
    path : str
        Path from which to get the data.
    df : panda dataframe
        Dataframe with the results from the optimization with Tune.

    Returns
    -------
    list
        Efficiency of each walk.
    list
        Parameters of the benchmark model that achieved the highest efficiency.
        Levy - list of the form [beta]
        Bi-exp. - list of the form [d_int, d_ext, p]

    """
    
    filter_ = df['mean_eff'] == max(df['mean_eff'])
    
    if model == 'powerlaw':
        par = [df[filter_]['config/beta'].values[0]]
    
    elif model == 'double_exp':
        par = [df[filter_]['config/d_int'].values[0], df[filter_]['config/d_ext'].values[0], df[filter_]['config/p'].values[0]]
        
    if path:
        effs = np.load(path+'efficiencies_'+str([np.round(p, 10) for p in par])+'.npy')
    else:
        effs = max(df['mean_eff'])

    return effs, par 
