# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/lib_nbs/11_environments_numba.ipynb.

# %% auto 0
__all__ = ['isBetween_c_Vec_numba', 'pareto_sample', 'rand_choice_nb', 'TargetEnv', 'reset_search_loop', 'ResetEnv_1D',
           'parallel_Reset1D_sharp', 'parallel_Reset1D_exp', 'ResetEnv_2D', 'parallel_Reset2D_sharp',
           'parallel_Reset2D_exp', 'parallel_Reset2D_policies', 'TurnResetEnv_2D', 'search_loop_turn_reset_sharp']

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 4
import numpy as np
import numba
from numba.experimental import jitclass
from numba import float64, bool_, prange, njit
import math
import random

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 8
@njit
def isBetween_c_Vec_numba(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.

        """
        if (a == b).all():
            return np.array([False]*c.shape[0])

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

        #exclude the targets that are beyond the step.
        for i3 in np.argwhere(dotproduct > squaredlengthba):
            mask[i3] = False
            
        return mask

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 14
@njit
def pareto_sample(alpha, xm, size=1):
    samples = np.zeros(size)
    for ii in range(size):
        u = random.random()  # Uniform random variable between 0 and 1
        x = xm / (u ** (1 / alpha))
        samples[ii] = x
    return samples

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 16
@njit
def rand_choice_nb(arr, prob):
    """
    :param arr: A 1D numpy array of values to sample from.
    :param prob: A 1D numpy array of probabilities for the given samples.
    :return: A random sample from the given array with a given probability.
    """
    return arr[np.searchsorted(np.cumsum(prob), np.random.random(), side="right")]

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 18
@jitclass([("target_positions", float64[:,:]) ,
           ("current_rewards", float64[:]) ,
           ("kicked", float64[:]) ,
           ("current_directions", float64[:]) ,
           ("positions", float64[:,:]),
           ("previous_pos", float64[:,:]),
           ("lc", float64[:,:]),
           ("mask", bool_[:]),
           ("first_encounter", float64[:,:])])
class TargetEnv():
    Nt : int
    L : float
    r : float
    lc : np.array
    agent_step : float
    num_agents : int
    destructive_targets : bool
    target_positions : np.ndarray
    current_rewards : np.array
    kicked : np.array
    current_directions : np.array
    positions : np.array
    previous_pos : np.array
    mask : np.array
    first_encounter : np.array
    lc_distribution : str
    
    
    def __init__(self,
                 Nt = 10, # Number of targets.
                 L = 1.3, #  Size of the (squared) world.
                 r = 1.5, # Radius with center the target position. It defines the area in which agent detects the target.
                 lc = np.array([[1.0],[1]]), # Cutoff length. Displacement away from target 
                 agent_step = 1, # Displacement of one step. The default is 1.
                 num_agents = 1, # Number of agents that forage at the same time. The default is 1. > 1 not fully implemented
                 destructive = False, # True if targets are destructive. The default is False.
                 lc_distribution = 'constant' # Distribution from where to sample l_c. Can be 'power_law', 'pareto' or something else. See comments self.check_encounter for explanations
                ):
        
        """
        Class defining the a Foraging environment with multiple targets and two actions: continue in 
        the same direction and turn by a random angle.
        """
        
        self.Nt = Nt
        self.L = L
        self.r = r
        self.lc = lc
        self.agent_step = agent_step 
        self.num_agents = num_agents
        self.destructive_targets = destructive
        self.lc_distribution = lc_distribution
        

        self.init_env()
        
    def init_env(self):
        """
        Environment initialization.
        """
        self.target_positions = np.random.rand(self.Nt, 2)*self.L
        
        #store who is/was rewarded
        self.current_rewards = np.zeros(self.num_agents)
        
        #signal whether agent has been kicked
        self.kicked = np.zeros(self.num_agents)
        
        #set positions and directions of the agents
        self.current_directions = np.random.rand(self.num_agents)*2*np.pi
        self.positions = np.random.rand(self.num_agents, 2)*self.L
        self.previous_pos = self.positions.copy()       

        

    def update_pos(self, 
                   change_direction, # Whether the agent decided to turn or not.
                   agent_index = 0 # Index of the given agent. The default is 0. This is only keeped for future devs
                  ):        
        """
        Updates information of the agent depending on its decision.            
        """
        
        # Save previous position to check if crossing happened
        self.previous_pos[agent_index] = self.positions[agent_index].copy()
        
        if change_direction:
            self.current_directions[agent_index] = random.uniform(0,1)*2*math.pi
        
        #Update position
        self.positions[agent_index][0] = self.positions[agent_index][0] + self.agent_step*np.cos(self.current_directions[agent_index])
        self.positions[agent_index][1] = self.positions[agent_index][1] + self.agent_step*np.sin(self.current_directions[agent_index])
        
       
    def check_encounter(self,
                       agent_index=0 # Index of the given agent. The default is 0. This is only keeped for future devs
                       ): # True if the agent found a target, else False
        """
        Checks whether the agent found a target, and updates the information accordingly.
        """       
        
        encounters = isBetween_c_Vec_numba(self.previous_pos[agent_index], self.positions[agent_index], self.target_positions, self.r)
        
        if sum(encounters) > 0: 
            
            #if there is more than 1 encounter, pick the closest to the agent.
            if sum(encounters) == 1:
                first_encounter = np.argwhere(encounters == True).flatten()
            else:
                # compute the distance from the previous position to each target            
                distance_previous_pos = np.sqrt((self.previous_pos[agent_index][0]- self.target_positions[:, 0])**2 + (self.previous_pos[agent_index][1] - self.target_positions[:, 1])**2)            
                
                # checking which encountered point is closer to previous position
                min_distance_masked = np.argmin(distance_previous_pos[encounters])
                first_encounter = np.argwhere(encounters == True)[min_distance_masked].flatten()
            if self.destructive_targets:
                self.target_positions[first_encounter] = np.random.rand(2)*self.L
            else:
                #----KICK----
                # If there was encounter, we reset direction and change position of particle to (pos target + lc)
                kick_direction = np.random.uniform(low = 0, high = 2*np.pi)
                for idx_first in first_encounter: # This is super weird!
                    if self.lc_distribution == 'power_law':
                        # when we have the power law, the first value of lc is considered to be the exponent.
                        # The following samples from a power law x^{-1-alpha} where alpha = self.lc.flatten()[0]                        
                        current_lc = (1-random.uniform(0,1))**(-1/self.lc.flatten()[0])

                    elif self.lc_distribution == 'pareto':
                        # Sampling from Pareto. Here alpha = self.lc.flatten()[0] and x_minim = self.lc.flatten()[0]
                        current_lc = pareto_sample(self.lc[0,0], self.lc[1,0])[0]
                    else:
                        # if lc has a single element, take that one as lc, if not sample
                        current_lc = self.lc.flatten()[0] if len(self.lc.flatten()) == 2 else rand_choice_nb(arr = self.lc[0], prob = self.lc[1])
                    self.positions[agent_index][0] = self.target_positions[idx_first, 0] + current_lc*np.cos(kick_direction)
                    self.positions[agent_index][1] = self.target_positions[idx_first, 1] + current_lc*np.sin(kick_direction)
                self.kicked[agent_index] = 1
                #------------
                
            #...and we add the information that this agent got to the target
            self.current_rewards[agent_index] = 1              
            return 1
        
        else: 
            self.kicked[agent_index] = 0
            self.current_rewards[agent_index] = 0
            return 0   
        
    def check_bc(self):
        """
        Updates position coordinates of agent agent_index to fulfill periodic boundary conditions.

        """
        agent_index=0
        self.positions[agent_index] = (self.positions[agent_index])%self.L
    

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 26
@njit
def reset_search_loop(T, # Number of steps 
                      reset_policy, # Reset policy
                      env # Environment
                      ):
    '''
    Loop that runs the reset environment with a given reset policy.
    '''
    
    rewards = 0
    tau = 0 # time since last reset
    
    for t in range(T):
        
        action = 0 if np.random.rand() > reset_policy[tau] else 1
        rew = env.update_pos(action = action)
        
        if rew == 1 or action == 1:
            tau = 0
        else:
            tau += 1
        
        rewards += rew
    
    return rewards


# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 28
@jitclass
class ResetEnv_1D():
    L : float
    D : float    
    position : float    
    
    def __init__(self,
                 L = 1.3,
                 D = 1.0,                    
                ):        
   
        self.L = L
        self.D = D
        self.position = 0
        
    def init_env(self):
        self.position = 0
    
    def update_pos(self, 
                   action # 0: continue walk, 1: reset to origin
                  ): # Reward = 1 if crossed L, else = 0
        
        if action == 0:
            self.position += np.random.randn()*np.sqrt(2*self.D)        
        else: self.position = 0
                
        if self.position >= self.L: 
            self.init_env()
            return 1
        else: return 0

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 30
@njit(parallel = True)
def parallel_Reset1D_sharp(T, resets, L, D):
    '''
    Runs the Reset 1D loop in parallel for different sharp resetting times.
    '''
    rews_reset = np.zeros_like(resets)
    
    for idxr in prange(len(resets)):
        
        env = ResetEnv_1D(L, D)        
        reset_policy = np.zeros(resets[idxr])
        reset_policy[resets[idxr]-1] = 1        
        
        rews_reset[idxr] = reset_search_loop(T = T, reset_policy = reset_policy, env = env)
    return rews_reset

@njit(parallel = True)
def parallel_Reset1D_exp(T, rates, L, D):
    '''
    Runs the Reset 1D loop in parallel for different exponential resetting rates.
    '''
    
    rews_rate = np.zeros_like(rates)
    for idxr in prange(len(rates)):
        
        env = ResetEnv_1D(L, D)        
        reset_policy = np.ones(T)*rates[idxr]
        
        rews_rate[idxr] = reset_search_loop(T = T, reset_policy = reset_policy, env = env)
    return rews_rate

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 34
@jitclass([("position", float64[:]),
           ("target_position", float64[:,:]),
           ("previous_pos", float64[:])
          ])
class ResetEnv_2D():
    D : float    
    dist_target :float
    r: float
    position : np.array    
    target_position : np.array
    previous_pos : np.array
    
    def __init__(self,
                 dist_target = 0.2, # Distance from init position and target
                 radius_target = 0.5, # Radius of the target
                 D = 1.0, # Diffusion coefficient of the walker              
                ):    
        

        self.D = D
        self.dist_target = dist_target
        self.r = radius_target        
        
        self.target_position = np.array([self.dist_target*np.cos(np.pi/4), self.dist_target*np.sin(np.pi/4)])[np.newaxis, :]       
        
        self.init_env()
        
    def init_env(self):
        self.position = np.array([0.0,0.0])
        self.previous_pos = self.position.copy()     
        
        
    
    def update_pos(self, 
                   action # 0: continue walk, 1: reset to origin
                  ): # Reward = 1 if encountered target, else = 0
        
        if action == 1:
            self.init_env()
            return 0
        
        elif action == 0:
            
            self.previous_pos = self.position.copy()            
            self.position += np.random.randn(2)*np.sqrt(2*self.D)
            
            # Checking encounter
            inside_target = np.linalg.norm(self.position-self.target_position) <= self.r
            crossed_target = isBetween_c_Vec_numba(self.previous_pos, self.position, self.target_position, self.r)
                        
            if inside_target or crossed_target:
                self.init_env()
                return 1
            else: 
                return 0
            
        

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 36
@njit(parallel = True)
def parallel_Reset2D_sharp(T, resets, dist_target, radius_target, D):
    rews_reset = np.zeros_like(resets)
    
    for idxr in prange(len(resets)): 
        
        env = ResetEnv_2D(dist_target, radius_target, D)        
        reset_policy = np.zeros(resets[idxr])
        reset_policy[resets[idxr]-1] = 1        
        
        rews_reset[idxr] = reset_search_loop(T = T, reset_policy = reset_policy, env = env)
    return rews_reset



@njit(parallel = True)
def parallel_Reset2D_exp(T, rates, dist_target, radius_target, D):
    
    rews_rate = np.zeros_like(rates)
    for idxr in prange(len(rates)):
        
        env = ResetEnv_2D(dist_target, radius_target, D)         
        reset_policy = np.ones(T)*rates[idxr]
        
        rews_rate[idxr] = reset_search_loop(T = T, reset_policy = reset_policy, env = env)
    return rews_rate

@njit(parallel = True)
def parallel_Reset2D_policies(T, reset_policies, dist_target, radius_target, D):
    
    rews_rate = np.zeros(reset_policies.shape[0])    
    
    for idx_policy in prange(reset_policies.shape[0]):    
        
        env = ResetEnv_2D(dist_target, radius_target, D)  
                
        rews_rate[idx_policy] = reset_search_loop(T = T, reset_policy = reset_policies[idx_policy], env = env)
        
    return rews_rate

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 39
@jitclass([("position", float64[:]),
           ("target_position", float64[:,:]),
           ("previous_pos", float64[:])
          ])
class TurnResetEnv_2D():
    agent_step : float    
    dist_target :float
    r: float
    position : np.array    
    target_position : np.array
    previous_pos : np.array
    current_direction : float
    
    def __init__(self,
                 dist_target = 0.2, # Distance from init position and target
                 radius_target = 0.5, # Radius of the target
                 agent_step = 1.0, # Diffusion coefficient of the walker              
                ):   
        
        '''
        Class defining a Foraging environment with a single target and three possible actions:

        - Continue in the same direction
        - Turn by a random angle
        - Reset to the origin

        The agent makes steps of constant length given by agent_step. 
        '''
        

        self.agent_step = agent_step
        self.dist_target = dist_target
        self.r = radius_target        
        
        self.target_position = np.array([self.dist_target*np.cos(np.pi/4), self.dist_target*np.sin(np.pi/4)])[np.newaxis, :]       
        
        self.init_env()
        
    def init_env(self):
        self.position = np.array([0.0,0.0])
        self.current_direction = np.random.rand()*2*np.pi
        self.previous_pos = self.position.copy()     
        
        
    
    def update_pos(self, 
                   change_direction, # If True, the agent changes direction by a random angle
                   reset # If True, the agent is reset to the origin
                   ):        
        """
        Updates position of the agent depending on its decision
        """
        
        if reset:
            self.init_env()
            return 0

        else:
            if change_direction:
                self.current_direction = np.random.rand()*2*np.pi

            self.position[0] += self.agent_step*np.cos(self.current_direction)
            self.position[1] += self.agent_step*np.sin(self.current_direction)

            # Checking encounter
            inside_target = np.linalg.norm(self.position-self.target_position) <= self.r
            crossed_target = isBetween_c_Vec_numba(self.previous_pos, self.position, self.target_position, self.r)
                        
            if inside_target or crossed_target:
                self.init_env()
                return 1
            else:                
                self.previous_pos = self.position.copy()
                return 0

# %% ../../../nbs/lib_nbs/11_environments_numba.ipynb 41
def search_loop_turn_reset_sharp(T, reset, turn, env):
    """
    Runs a search loop of T steps. There is a single counter that works as follows:

    - Starts at 0
    - For each turn or continue action gets +1
    - If reset or reach the target is set to 0

    """
    
    rewards = 0
    counter = 0 

    env.init_env()
    
    for t in range(T):        
        counter += 1
        # Reset
        if counter == reset:
            rew = env.update_pos(False, # change direction
                                 True   # reset
                                )           
            counter = 0
        
        # Turn
        elif counter == turn:
            rew = env.update_pos(True, # change direction
                                 False # reset
                                )
            
            
        else:
            rew = env.update_pos(False, # change direction
                                 False # reset
                                )
        if rew == 1:
            counter = 0
            
        rewards += rew
        
    return rewards
