# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/models_phenom.ipynb (unless otherwise specified).

__all__ = ['models_phenom', 'models_phenom', 'models_phenom', 'models_phenom', 'models_phenom', 'models_phenom',
           'models_phenom', 'models_phenom', 'models_phenom', 'models_phenom', 'models_phenom']

# Cell
import numpy as np
import fbm
import matplotlib.pyplot as plt
import warnings

try:
    __IPYTHON__
    from tqdm.notebook import tqdm
except NameError:
    from tqdm import tqdm


# Cell
class models_phenom():
    def __init__(self):
        '''Constructor of the class'''

# Cell
class models_phenom(models_phenom):

    @staticmethod
    def _single_state_traj(T = 200,
                          D = 1,
                          alpha = 1,
                          L = None):

        # Trajectory
        dispx, dispy = fbm.fgn(n = T, hurst = alpha/2), fbm.fgn(n = T, hurst = alpha/2)
        dispx, dispy = np.sqrt(D)*(dispx/np.std(dispx)), np.sqrt(D)*(dispy/np.std(dispy))
        # labels
        labels = np.vstack((np.ones(T)*alpha, np.ones(T)*D)).transpose()

        # If there are no boundaries
        if not L:
            posx, posy = np.cumsum(dispx) - dispx[0], np.cumsum(dispy) - dispy[0]

            return np.vstack((posx, posy)).transpose(), labels

        # If there are, apply reflecting boundary conditions
        else:
            pos = np.zeros((T, 2))

            # Initialize the particle in a random position of the box
            pos[0, :] = np.random.rand(2)*L
            for t in range(1, T):
                pos[t, :] = [pos[t-1, 0]+dispx[t], pos[t-1, 1]+dispy[t]]


                # Reflecting boundary conditions
                while np.max(pos[t, :])>L or np.min(pos[t, :])< 0:
                    pos[t, pos[t, :] > L] = pos[t, pos[t, :] > L] - 2*(pos[t, pos[t, :] > L] - L)
                    pos[t, pos[t, :] < 0] = - pos[t, pos[t, :] < 0]

            return pos, labels

# Cell
class models_phenom(models_phenom):


    def single_state(self,
                     N = 10,
                     T = 200,
                     D = 1,
                     alpha = 1,
                     L = None):

        data = np.zeros((T, N, 2))
        labels = np.zeros((T, N, 2))

        for n in range(N):

            # Get trajectory from single traj function
            pos, lab = self._single_state_traj(T = T,
                                   D = D,
                                   alpha = alpha,
                                   L = L)
            data[:, n, :] = pos
            labels[:, n, :] = lab

        return data, labels

# Cell
class models_phenom(models_phenom):

    @staticmethod
    def _multiple_states_traj(T = 200,
                    M = np.array([[0.9 , 0.1],[0.1 ,0.9]]),
                    Ds = np.array([10.,0]),
                    alphas = np.array([1., 1]),
                    L = None):

        # transform lists to numpy if needed
        if isinstance(M, list):
            M = np.array(M)
        if isinstance(Ds, list):
            Ds = np.array(Ds)
        if isinstance(alphas, list):
            alphas = np.array(alphas)


        pos = np.zeros((T, 2))
        if L: pos[0,:] = np.random.rand(2)*L

        # Diffusing state of the particle
        state = np.zeros(T).astype(int)
        state[0] = np.random.randint(M.shape[0])
        # Output labels
        labels = np.zeros((T, 2))
        labels[0, 0] = alphas[state[0]]
        labels[0, 1] = Ds[state[0]]

        # Trajectory
        dispx, dispy = fbm.fgn(n = T, hurst = alphas[state[0]]/2), fbm.fgn(n = T, hurst = alphas[state[0]]/2)
        dispx, dispy = np.sqrt(Ds[state[0]])*(dispx/np.std(dispx)), np.sqrt(Ds[state[0]])*(dispy/np.std(dispy))
        disp_t = 0


        for t in range(1, T):

            pos[t, :] = [pos[t-1, 0]+dispx[disp_t], pos[t-1, 1]+dispy[disp_t]]
            disp_t += 1
            # at each time, check new state
            state[t] = np.random.choice(np.arange(M.shape[0]), p = M[state[t-1], :])


            if state[t] != state[t-1]:
                # Recalculate new displacements for next steps
                dispx, dispy = [fbm.fgn(n = T-t, hurst = alphas[state[t]]/2),
                                fbm.fgn(n = T-t, hurst = alphas[state[t]]/2)]

                if len(dispx) > 1:
                    dispx, dispy = [np.sqrt(Ds[state[t]])*(dispx/np.std(dispx)),
                                    np.sqrt(Ds[state[t]])*(dispy/np.std(dispy))]
                else:
                    dispx, dispy = [np.sqrt(Ds[state[t]])*np.random.randn(),
                                    np.sqrt(Ds[state[t]])*np.random.randn()]
                disp_t = 0

            if L is not None:
                # Reflecting boundary conditions
                while np.max(pos[t, :])>L or np.min(pos[t, :])< 0:
                    pos[t, pos[t, :] > L] = pos[t, pos[t, :] > L] - 2*(pos[t, pos[t, :] > L] - L)
                    pos[t, pos[t, :] < 0] = - pos[t, pos[t, :] < 0]


            labels[t, 0] = alphas[state[t]]
            labels[t, 1] = Ds[state[t]]

             # Update labels


        return pos, labels



# Cell
class models_phenom(models_phenom):
    def multi_state(self,
                    N = 10,
                    T = 200,
                    M = np.array([[0.9 , 0.1],[0.1 ,0.9]]),
                    Ds = np.array([10.,0]),
                    alphas = np.array([0.5, 1]),
                    L = None):

        data = np.zeros((T, N, 2))
        labels = np.zeros((T, N, 2))

        for n in range(N):

            # Get trajectory from single traj function
            pos, lab = self._multiple_states_traj(T = T,
                                                   L = L,
                                                   M = M,
                                                   alphas = alphas,
                                                   Ds = Ds)
            data[:, n, :] = pos
            labels[:, n, :] = lab

        return data, labels

# Cell
class models_phenom(models_phenom):
    @staticmethod
    def _get_distance(x):
        M = np.reshape(np.repeat(x[ :, :], x.shape[0], axis = 0), (x.shape[0], x.shape[0], 2))
        Mtrans = M.transpose(1,0,2)
        distance = np.sqrt(np.square(M[:,:, 0]-Mtrans[:,:, 0])
                         + np.square(M[:,:, 1]-Mtrans[:,:, 1]))
        return distance

# Cell
class models_phenom(models_phenom):
    @staticmethod
    def _make_escape(Pu, label, diff_state):

        # if unbinding probability is zero
        if Pu == 0:
            return label, diff_state

        label = label.copy()
        diff_state = diff_state.copy()

        label_dimers = np.unique(label[np.argwhere(diff_state == 1)])

        for l in label_dimers:

            if np.random.rand() < Pu:
                # give new label to escaping particles
                diff_state[label == l] = 0
                label[label == l] = np.max(label)+np.arange(2)+1

        return label, diff_state

# Cell
class models_phenom(models_phenom):
    @staticmethod
    def _make_condensates(Pb, label, diff_state, r, distance, max_label):

        label = label.copy()
        diff_state = diff_state.copy()

        # Keeping track of the ones that will dimerize
        already_dimer = []

        for n, l in enumerate(label):

            # Consider conditions in which particles do not dimerize
            if n in already_dimer or diff_state[n] == 1 or l > max_label:
                continue

            # Extract distances to current particle
            distance_to_current = distance[n,:]
            distance_to_current[n] == 0
            close_particles = np.argwhere((distance_to_current < 2*r) & (distance_to_current > 0)).flatten()

            # Loop over all posible dimerizing candidates
            for chosen in close_particles:

                # Consider conditions in which particles do not dimerize
                if chosen in already_dimer or diff_state[chosen] == 1 or label[chosen] > max_label:
                    continue

                # Draw coin to see if particle dimerizes
                if np.random.rand() < Pb:
                    # Add dimerized particles to the new dimer counter
                    already_dimer.append(chosen)
                    already_dimer.append(n)
                    # Update their diffusive state
                    diff_state[n] = 1
                    diff_state[chosen] = 1

                    # dimerize particles
                    label[chosen] = l

                    # if one particles dimers, not more clustering!
                    break

        return label, diff_state

# Cell
class models_phenom(models_phenom):
    def dimerization(self,
                     N = 10,
                     T = 100,
                     L = 5,
                     r = 1,
                     Pu = 0.1, # Unbinding probability
                     Pb = 0.01, # Binding probability
                     Ds = np.array([1, 0]), # Diffusion coefficients of two states
                     alphas = np.array([1, 1]) # Anomalous exponents for two states
                     ):

        # transform lists to numpy if needed
        if isinstance(Ds, list):
            Ds = np.array(Ds)
        if isinstance(alphas, list):
            alphas = np.array(alphas)

        # Info to save
        pos = np.zeros((T, N, 2)) # position over time
        label = np.zeros((T, N)).astype(int)
        diff_state = np.zeros((T, N)).astype(int)
        output_label = np.zeros((T, N, 2))

        # Init position, labels
        pos[0, :, :] = np.random.rand(N, 2)*L
        label[0, :] = np.arange(pos.shape[1])
        output_label[0, :, 0] = alphas[diff_state[0, :]]
        output_label[0, :, 1] = Ds[diff_state[0, :]]

        # initial displacements (all free particles)
        disps = np.zeros((T, N, 2))
        for n in range(N):
            dispx, dispy = [fbm.fgn(n = T, hurst = alphas[0]/2),
                            fbm.fgn(n = T, hurst = alphas[0]/2)]

            disps[:, n, 0] = np.sqrt(Ds[0])*(dispx/np.std(dispx))
            disps[:, n, 1] = np.sqrt(Ds[0])*(dispy/np.std(dispy))


        for t in tqdm(range(1, T)):

            # Find max label to account later for escaped
            max_label = np.max(label[t-1, :])

            # Make particles escape
            label[t, :], diff_state[t, :] = self._make_escape(Pu,
                                                              label[t-1, :],
                                                              diff_state[t-1, :])

            lab, diff = label[t, :].copy(), diff_state[t, :].copy()

            # get distance + increasing it for escaped to avoid reclustering
            distance = self._get_distance(pos[t-1, :, :])

            # Merge particles in condensates
            label[t, :], diff_state[t, :] = self._make_condensates(Pb,
                                                                 label[t, :],
                                                                 diff_state[t, :],
                                                                 r, distance, max_label)

            # Find particles which changed state
            label_changed, counts = np.unique(label[t, np.not_equal(diff_state[t-1,:], diff_state[t,:])],
                                              return_counts = True)

            # Calculate new displacements for particles which changed state
            for l, count in zip(label_changed, counts):

                index = int(np.argwhere(label[t,:] == l)[0])

                current_alpha = alphas[diff_state[t, index]]
                current_D = Ds[diff_state[t, index]]

                dispx, dispy = [fbm.fgn(n = T-t, hurst = current_alpha/2),
                                fbm.fgn(n = T-t, hurst = current_alpha/2)]

                if len(dispx) > 1:
                    disps[t:, label[t, :] == l, 0] = np.tile(current_D*(dispx/np.std(dispx)), (count, 1)).transpose()
                    disps[t:, label[t, :] == l, 1] = np.tile(current_D*(dispy/np.std(dispy)), (count, 1)).transpose()
                else:
                    disps[t:, label[t, :] == l, 0] = current_D*np.random.randn(1, count)
                    disps[t:, label[t, :] == l, 1] = current_D*np.random.randn(1, count)

            # Update position
            pos[t, :, :] = pos[t-1,:,:]+disps[t, :, :]
            # Consider boundary conditions
            if L is not None:
                while np.max(pos[t,:, :])>L or np.min(pos[t,:, :])< 0:
                    pos[t, pos[t,:, :] > L] = pos[t, pos[t,:, :] > L] - 2*(pos[t, pos[t,:, :] > L] - L)
                    pos[t, pos[t,:, :] < 0] = - pos[t, pos[t,:, :] < 0]

            # Update output labels
            output_label[t, :, 0] = alphas[diff_state[t, :]]
            output_label[t, :, 1] = Ds[diff_state[t, :]]

        return pos, output_label


# Cell
class models_phenom(models_phenom):
    @staticmethod
    def _update_bound(mask, # current bind vector
                         N, # number of particles
                         pos, # position of particles
                         Nt, # number of traps
                         traps_pos, # position of traps
                         Pb, # binding probability
                         Pu, # unbinding probability
                         r, # trap radius
                        ):

        # from the ones that are bound, get the ones that unbind. These will be descarted for binding in same time step
        mask_new_free = np.array(1-(np.random.rand(N) < Pu)*mask).astype(bool)

        # calculate the distance between traps and particles
        d = models_phenom._get_distance(np.vstack((traps_pos, pos)))[Nt:, :Nt]
        mask_close = (d < r).sum(1).astype(bool)

        # get mask for binding
        mask_new_bind = np.random.rand(N) < Pb

        # update the bound vector with the previous conditions:
        # first, the ones that unbind
        mask *= mask_new_free
        # then, the ones that are close + bind. Mask_new_free is added to avoid binding
        # of the ones that just unbound
        mask += mask_close*mask_new_bind*mask_new_free

        return mask

# Cell
class models_phenom(models_phenom):

    def immobile_traps(self,
                       N = 10,
                        T = 100,
                        L = 5,
                        r = 1,
                        Pu = 0.1, # Unbinding probability
                        Pb = 0.01, # Binding probability
                        D = 1, # Diffusion coefficients of moving states
                        alpha = 1, # Anomalous exponents for two states
                        Nt = 10,
                        traps_pos = None
                        ):

        # Info to output
        pos = np.zeros((T, N, 2)) # position over time
        output_label = np.zeros((T, N, 2))

        disps = np.zeros((T, N, 2))
        diff_state = np.zeros((T, N)).astype(int)
        mask_bound = diff_state[0, :].astype(bool)

        # Init position, labels
        pos[0, :, :] = np.random.rand(N, 2)*L
        output_label[0, :, 0] = alpha
        output_label[0, :, 1] = D


        # Traps positions
        if traps_pos is None:
            traps_pos = np.random.rand(Nt, 2)*L


        for n in range(N):
            dispx, dispy = [fbm.fgn(n = T, hurst = alpha/2),
                            fbm.fgn(n = T, hurst = alpha/2)]

            disps[:, n, 0] = np.sqrt(D)*(dispx/np.std(dispx))
            disps[:, n, 1] = np.sqrt(D)*(dispy/np.std(dispy))


        for t in tqdm(range(1, T)):

            mask_bound = self._update_bound(mask = mask_bound, # current bind vector
                                         N = N, # number of particles
                                         pos = pos[t-1, :, :], # position of particles
                                         Nt = Nt, # number of traps
                                         traps_pos = traps_pos, # position of traps
                                         Pb = Pb, # binding probability
                                         Pu = Pu, # unbinding probability
                                         r = r, # trap radius
                                         )
            # Update the diffusive state
            diff_state[t,:] = mask_bound

            # Regenerate trajectories for untrapped particles
            untrapped = np.argwhere((diff_state[t,:] - diff_state[t-1,:]) == -1).flatten()
            for un_part in untrapped:
                    if T-t > 1:
                        # Recalculate new displacements for next steps
                        dispx, dispy = [fbm.fgn(n = T-t, hurst = alpha/2),
                                        fbm.fgn(n = T-t, hurst = alpha/2)]

                        disps[t:, un_part, 0] = np.sqrt(D)*(dispx/np.std(dispx))
                        disps[t:, un_part, 1] = np.sqrt(D)*(dispy/np.std(dispy))

                    else:
                        disps[t:, un_part, 0] = np.sqrt(D)*np.random.randn()
                        disps[t:, un_part, 1] = np.sqrt(D)*np.random.randn()

            # Update the position
            pos[t, :, :] = pos[t-1, :, :] + (1-mask_bound).reshape(N,1)*disps[t, :, :]

            # Update labels
            output_label[t, :, 0] = alpha*(1-mask_bound)
            output_label[t, :, 1] = D*(1-mask_bound)

            # Consider boundary conditions
            if L is not None:
                while np.max(pos[t,:, :])>L or np.min(pos[t,:, :])< 0:
                    pos[t, pos[t,:, :] > L] = pos[t, pos[t,:, :] > L] - 2*(pos[t, pos[t,:, :] > L] - L)
                    pos[t, pos[t,:, :] < 0] = - pos[t, pos[t,:, :] < 0]

        return pos, output_label