# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/lib_nbs/05_mfpt.ipynb.

# %% auto 0
__all__ = ['constant_velocity_generator', 'mfpt_rw', 'mfpt_informed_rw', 'rw_generator', 'exp_time_generator', 'Biexp',
           'biexp_time_generator', 'constant_velocity_generator_2D', 'mfpt_rw_2D', 'mfpt_informed_rw_2D']

# %% ../nbs/lib_nbs/05_mfpt.ipynb 3
import numpy as np
from tqdm.notebook import tqdm, trange
from tqdm import tqdm as tqdm_term
from typing import Callable

from .analytics import pdf_powerlaw, pdf_discrete_sample

# %% ../nbs/lib_nbs/05_mfpt.ipynb 4
# This package is not included in the requirements and should be installed by hand
from joblib import Parallel, delayed

# %% ../nbs/lib_nbs/05_mfpt.ipynb 7
def constant_velocity_generator(N, # Number of trajectories
                                T, # Length of trajectories
                                time_sampler, # Sampler of time of constant velocity
                                velocity = 1, # Velocity
                                **sample_args # Optional arguments for the sampler
                               ):
    ''' Given a sampler for length of time steps, generates a trajectory considering a constant
    velocity in the sampled times. After each time step, we sample a new direction. '''
    
    trajs = np.zeros((N, T))
    
    for idxN in (range(N)):        
        
        times = time_sampler(num_samples = T+1, **sample_args).astype(int)
        
        max_t = np.argwhere(np.cumsum(times) > T+1).flatten()
        if len(max_t) > 0:
            max_t = max_t[0]
        else: max_t = T+1
        
        traj = np.array([0])
        cumsum = 0
        for time in times[:max_t+1]:
            segment = (2*np.random.randint(0, 2)-1)*np.arange(1, min(time+1, T+1))+traj[-1]
            
            traj = np.append(traj, segment)
            cumsum += len(segment)
            if cumsum > T:
                break

        trajs[idxN] = traj[:T]
        
    return trajs
    

# %% ../nbs/lib_nbs/05_mfpt.ipynb 12
from typing import Callable

def mfpt_rw(N:int, # Number of trajectories
            T:int, # Lengths of trajectories
            x0:float, # Starting point of walk
            Ls:list, # Boundary position
            traj_generator:Callable,  # Function generating the walk (must start in x0 = zero)
            max_loop = 5,  # Max number of while loop if some walks do not reach the boundary
            save = None,
            **args_generator # Arguments of the trajectory generator
           ): # MFPT vs. L
    ''' Calculates MFPT to boundaries at 0 and L starting in x0 '''
    
    trajs = traj_generator(N,T,**args_generator) + x0
    
    if save is not None:
        statistics_fpt = np.zeros((len(Ls), N))

    mfpt = []
    for idxL, L in tqdm_term(enumerate(Ls)):   
        first_passage = np.argmax((trajs < 0) | (trajs > L), axis = 1) 

        # Enlarging the time of those who did not hit the boundaries
        n_loop = 0
        non_passed = np.argwhere(first_passage == 0).flatten()
        last_non_passed = trajs[non_passed, -1]
        while len(non_passed) > 0 and n_loop < max_loop:
            
            trajs_upd = traj_generator(len(non_passed),T,**args_generator) + last_non_passed[..., np.newaxis]
            
            # Finding the FPT of the new trajectories
            first_upd = np.argmax((trajs_upd < 0) | (trajs_upd > L), axis = 1)            
            # Adding these to the main FPT array, considering that loop number
            first_passage[non_passed] = (first_upd + (n_loop+1)*T)*first_upd.astype(bool)

            # Checking the ones that still did not pass
            non_passed = np.argwhere(first_passage == 0).flatten()
            # updating the last element of the non passed trajectories
            upd_non_passed = np.argwhere(first_upd == 0).flatten()
            if len(non_passed) != len(upd_non_passed): print('ERROR')
            last_non_passed = trajs_upd[upd_non_passed, -1]
            # updating loop num
            n_loop += 1      
                                  
        if save is not None:
            statistics_fpt[idxL] = first_passage.copy()
                                

        # Calculating mean considering that some may not have passed              
        first_passage = first_passage[first_passage != 0].mean() 
        # Adding this MFPT for current L
        mfpt.append(first_passage)
        
        
        
    if save is not None:
        np.save(save+'.npy', mfpt)
        np.save(save+'_stats.npy', statistics_fpt)
        
    return mfpt

# %% ../nbs/lib_nbs/05_mfpt.ipynb 14
def mfpt_informed_rw(N:int, # Number of trajectories
            T:int, # Lengths of trajectories
            x0:float, # Starting point of walk
            Ls:list, # Boundary position
            traj_generator:Callable,  # Function generating the walk (must start in x0 = zero)
            max_loop = 5,  # Max number of while loop if some walks do not reach the boundary
            time_sampler = None,
            save = None,
            **args_generator # Optional                     
           ): # MFPT vs. L
    ''' Calculates MFPT to boundaries at 0 and L starting in x0 with a traj generator informed on the scales of the sysmte (x0 and L)'''
    
    if save is not None:
        statistics_fpt = np.zeros((len(Ls), N))

    mfpt = []
    for idxL, L in tqdm_term(enumerate(Ls)):  
        
        trajs = traj_generator(N, T, time_sampler = time_sampler, x0 = x0, L = L, informed = True)
        
        first_passage = np.argmax((trajs < 0) | (trajs > L), axis = 1) 

        # Enlarging the time of those who did not hit the boundaries
        n_loop = 0
        non_passed = np.argwhere(first_passage == 0).flatten()
        while len(non_passed) > 0 and n_loop < max_loop:
            last_non_passed = trajs[non_passed, -1]
            trajs_upd = traj_generator(len(non_passed),T, 
                                       time_sampler = time_sampler, x0 = x0, L = L, 
                                       informed = True) + last_non_passed[..., np.newaxis]

            # updating trajectories with new positions. This makes that last_non_passed is always correct
            trajs[non_passed, -1] = trajs_upd[:, -1]
            # Finding the FPT of the new trajectories
            first_upd = np.argmax((trajs_upd < 0) | (trajs_upd > L), axis = 1)            
            # Adding these to the main FPT array, considering that loop number
            first_passage[non_passed] = (first_upd + (n_loop+1)*T)*first_upd.astype(bool)

            # Checking the ones that still did not pass
            non_passed = np.argwhere(first_passage == 0).flatten()
            # updating loop num
            n_loop += 1      
                                  
        if save is not None:
            statistics_fpt[idxL] = first_passage.copy()
                                

        # Calculating mean considering that some may not have passed              
        first_passage = first_passage[first_passage != 0].mean() 
        # Adding this MFPT for current L
        mfpt.append(first_passage)
        
        
        
    if save is not None:
        np.save(save+'.npy', mfpt)
        np.save(save+'_stats.npy', statistics_fpt)
        
    return mfpt

# %% ../nbs/lib_nbs/05_mfpt.ipynb 17
def rw_generator(N, T):
    trajs = 2*np.random.randint(0, 2, size = (N,T))-1
    trajs[:,0] = 0
    trajs = np.cumsum(trajs, axis = 1)
    return trajs

# %% ../nbs/lib_nbs/05_mfpt.ipynb 21
def exp_time_generator(num_samples):
    return np.random.exponential(size = num_samples)+1

# %% ../nbs/lib_nbs/05_mfpt.ipynb 25
from scipy.stats.sampling import NumericalInverseHermite

class Biexp():
    def __init__(self, informed = False, **args):
        
        if not informed:
            self.w1 = args['w1']
            self.d1 = args['d1']
            self.d2 = args['d2']
        else:
            self.w1 = 2*args['x0']/args['L']
            self.d1 = 3*args['L']
            self.d2 = args['x0']        
            
        
    def pdf(self, x):
        return (self.w1/self.d1)*np.exp(-x/self.d1) + ((1 - self.w1)/ self.d2)*np.exp(-x/self.d2)
    
    def cdf(self, x):
        return 1 + (self.w1-1)*np.exp(-x/self.d2) - self.w1*np.exp(-x/self.d1)
    
    def dpdf(self, x):
        return (self.w1/self.d1**2)*np.exp(-x/self.d1) + ((1 - self.w1)/ self.d2**2)*np.exp(-x/self.d2)
    
    def sample(self, num_rn):
        return NumericalInverseHermite(self, domain=(0, np.inf)).rvs(int(num_rn))

# %% ../nbs/lib_nbs/05_mfpt.ipynb 26
def biexp_time_generator(num_samples, **sample_kwargs):
    sampler = Biexp(**sample_kwargs)
    return sampler.sample(num_samples)+1

# %% ../nbs/lib_nbs/05_mfpt.ipynb 62
def constant_velocity_generator_2D(N, # Number of trajectories
                                   T, # Length of trajectories
                                   time_sampler, # Sampler of time of constant velocity
                                   velocity = 1, # Velocity
                                   **sample_args # Optional arguments for the sampler
                               ):
    ''' Given a sampler for length of time steps, generates a 2D trajectory considering a constant
    velocity in the sampled times. After each time step, we sample a new direction. '''
    
    trajs_x = np.zeros((N, T))
    trajs_y = np.zeros((N, T))
    
    for idxN in (range(N)):        
        
        times = time_sampler(num_samples = T+1, **sample_args).astype(int)
        
        # If we have too many times (e.g. because we sampled big times),
        # we cut the times vector with max_t
        max_t = np.argwhere(np.cumsum(times) > T+1).flatten()
        if len(max_t) > 0:
            max_t = max_t[0]
        else: max_t = T+1
        
        traj_x = np.array([0])
        traj_y = np.array([0])
        cumsum = 0
        for time in times[:max_t+1]:
            angle = np.random.rand()*2*np.pi
            
            segment_x = np.cos(angle)*np.arange(1, min(time+1, T+1))+traj_x[-1]
            segment_y = np.sin(angle)*np.arange(1, min(time+1, T+1))+traj_y[-1]
            
            traj_x = np.append(traj_x, segment_x)
            traj_y = np.append(traj_y, segment_y)
            
            cumsum += len(segment_x)
            if cumsum > T:
                break

        trajs_x[idxN] = traj_x[:T]
        trajs_y[idxN] = traj_y[:T]        
    
    return trajs_x, trajs_y
    

# %% ../nbs/lib_nbs/05_mfpt.ipynb 66
from typing import Callable

def mfpt_rw_2D(N:int, # Number of trajectories
               T:int, # Lengths of trajectories
               x0:list, # Starting point of walk (in 2d)
               Ls:list, # Boundary position
               traj_generator:Callable,  # Function generating the walk (must start in x0 = zero)
               max_loop = 5,  # Max number of while loop if some walks do not reach the boundary
               **args_generator # Arguments of the trajectory generator
           ): # MFPT vs. L
    ''' Calculates MFPT to boundaries at 0 and L starting in x0 in 2D'''



    # Debug
    # N = 5; T = 50; x0 = [3,3]; Ls = [99]; max_loop = 40

    trajs_x, trajs_y = traj_generator(N, T, **args_generator)
    
    trajs_x += x0[0]
    trajs_y += x0[1]

    mfpt = []
    for L in tqdm_term(Ls):   

        first_passage_x = np.argmax((trajs_x < -L**2) | (trajs_x > L), axis = 1).astype(float)
        first_passage_y = np.argmax((trajs_y < -L**2) | (trajs_y > L), axis = 1).astype(float)

        # Needed transformation to correctly find the minimum considering the presence of zeros
        first_passage_x[first_passage_x == 0] = np.nan
        first_passage_y[first_passage_y == 0] = np.nan    
        first_passage = np.nanmin(np.vstack((first_passage_x, first_passage_y)), axis = 0)
        first_passage[np.isnan(first_passage)] = 0

        # Enlarging the time of those who did not hit the boundaries
        n_loop = 0
        non_passed = np.argwhere(first_passage == 0).flatten()
        last_non_passed = (trajs_x[non_passed, -1], trajs_x[non_passed, -1])    
        while len(non_passed) > 0 and n_loop < max_loop:


            # print('inside')
            trajs_upd_x = rw_generator(len(non_passed),T) + last_non_passed[0][..., np.newaxis]
            trajs_upd_y = rw_generator(len(non_passed),T) + last_non_passed[1][..., np.newaxis]


            # Finding the FPT of the new trajectories
            first_upd_x = np.argmax((trajs_upd_x < -L**2) | (trajs_upd_x > L), axis = 1).astype(float)
            first_upd_y = np.argmax((trajs_upd_y < -L**2) | (trajs_upd_y > L), axis = 1).astype(float)

            # Needed transformation to correctly find the minimum considering the presence of zeros
            first_upd_x[first_upd_x == 0] = np.nan
            first_upd_y[first_upd_y == 0] = np.nan    
            first_upd = np.nanmin(np.vstack((first_upd_x, first_upd_y)), axis = 0)
            first_upd[np.isnan(first_upd)] = 0

            # Adding these to the main FPT array, considering that loop number
            first_passage[non_passed] = (first_upd + (n_loop+1)*T)*first_upd.astype(bool)

            # Checking the ones that still did not pass
            non_passed = np.argwhere(first_passage == 0).flatten()
            # updating the last element of the non passed trajectories
            upd_non_passed = np.argwhere(first_upd == 0).flatten()
            if len(non_passed) != len(upd_non_passed): print('ERROR')
            last_non_passed = (trajs_upd_x[upd_non_passed, -1], trajs_upd_y[upd_non_passed, -1])

            n_loop += 1        

        # Calculating mean considering that some may not have passed              
        first_passage = first_passage[first_passage != 0].mean() 
        if len(non_passed) > 0: print(round(len(non_passed)/len(non_passed),2))
        # Adding this MFPT for current L
        mfpt.append(first_passage)
        
    return mfpt
    

# %% ../nbs/lib_nbs/05_mfpt.ipynb 80
from typing import Callable

def mfpt_informed_rw_2D(N:int, # Number of trajectories
                        T:int, # Lengths of trajectories
                        x0:list, # Starting point of walk (in 2d)
                        Ls:list, # Boundary position
                        traj_generator:Callable,  # Function generating the walk (must start in x0 = zero)
                        max_loop = 5,  # Max number of while loop if some walks do not reach the boundary
                        **args_generator # Arguments of the trajectory generator
                       ): # MFPT vs. L
    ''' Calculates MFPT to boundaries at 0 and L starting in x0 in 2D'''



    # Debug
    # N = 5; T = 50; x0 = [3,3]; Ls = [99]; max_loop = 40


    mfpt = []
    for L in tqdm_term(Ls):   
        
        
        trajs_x, trajs_y = traj_generator(N, T, time_sampler = time_sampler, x0 = x0, L = L, informed = True)        
        trajs_x += x0[0]
        trajs_y += x0[1]

        first_passage_x = np.argmax((trajs_x < 0) | (trajs_x > L), axis = 1).astype(float)
        first_passage_y = np.argmax((trajs_y < 0) | (trajs_y > L), axis = 1).astype(float)

        # Needed transformation to correctly find the minimum considering the presence of zeros
        first_passage_x[first_passage_x == 0] = np.nan
        first_passage_y[first_passage_y == 0] = np.nan    
        first_passage = np.nanmin(np.vstack((first_passage_x, first_passage_y)), axis = 0)
        first_passage[np.isnan(first_passage)] = 0

        # Enlarging the time of those who did not hit the boundaries
        n_loop = 0
        non_passed = np.argwhere(first_passage == 0).flatten()
        last_non_passed = (trajs_x[non_passed, -1], trajs_x[non_passed, -1])    
        while len(non_passed) > 0 and n_loop < max_loop:


            # print('inside')
            trajs_upd_x = rw_generator(len(non_passed),T) + last_non_passed[0][..., np.newaxis]
            trajs_upd_y = rw_generator(len(non_passed),T) + last_non_passed[1][..., np.newaxis]


            # Finding the FPT of the new trajectories
            first_upd_x = np.argmax((trajs_upd_x < 0) | (trajs_upd_x > L), axis = 1).astype(float)
            first_upd_y = np.argmax((trajs_upd_y < 0) | (trajs_upd_y > L), axis = 1).astype(float)

            # Needed transformation to correctly find the minimum considering the presence of zeros
            first_upd_x[first_upd_x == 0] = np.nan
            first_upd_y[first_upd_y == 0] = np.nan    
            first_upd = np.nanmin(np.vstack((first_upd_x, first_upd_y)), axis = 0)
            first_upd[np.isnan(first_upd)] = 0

            # Adding these to the main FPT array, considering that loop number
            first_passage[non_passed] = (first_upd + (n_loop+1)*T)*first_upd.astype(bool)

            # Checking the ones that still did not pass
            non_passed = np.argwhere(first_passage == 0).flatten()
            # updating the last element of the non passed trajectories
            upd_non_passed = np.argwhere(first_upd == 0).flatten()
            if len(non_passed) != len(upd_non_passed): print('ERROR')
            last_non_passed = (trajs_upd_x[upd_non_passed, -1], trajs_upd_y[upd_non_passed, -1])

            n_loop += 1        

        # Calculating mean considering that some may not have passed              
        first_passage = first_passage[first_passage != 0].mean() 
        # Adding this MFPT for current L
        mfpt.append(first_passage)
        
    return mfpt
    
