import numpy as np          #import libraries
import matplotlib.pyplot as plt 
from matplotlib import animation 

plt.rcParams.update({"text.usetex": True})  #use latex




do_animations = False         #set to true if you want to see cool stuff idk 



"""
--------------------------------------------------- Q1 (a)
"""




        #define a square grid, of side length 2*L
L = 50          
axis = np.arange(-L, L+1)       #create an axis and a grid 
x, y = np.meshgrid(axis, axis)      


def get_cell_action(cell, x_axis, y_axis):          #function to calculate all the possible moves of a particle     
    """
    Returns list of possible accessible cells for an entry. Checks boundaries.

    Cell: array of indices.
    """
    i = cell[0]     #get cell indices
    j = cell[1]
    N = len(x_axis)-1       #get axis lengths
    M = len(y_axis)-1

    if (i<0) or (i>=N+1) or (j<0) or (j>=M+1):      #if index out of bounds, raise error 
        raise ValueError('i or j must be located within grid range!')
    
    xout = []       #blank return arrays 
    yout = []

        #boundaries
    if i == N:
        xout.append(N-1)    #move west

    if i == 0:
        xout.append(i+1)    #move east 

        #interior
    elif 0<i<N:
        xout.append(i-1)
        xout.append(i+1)


        #boundaries 
    if j == 0:
        yout.append(j+1)    #move north
    
    if j == M:
        yout.append(M-1)    #move south 

        #interior
    elif 0<j<M:
        yout.append(j-1)
        yout.append(j+1)

    return xout, yout       #return all the x and y possible moves




def next_cell(prev, x, y):      #function to propagate the particle to the next cell, randomized, one step
    """
    Forward-time stepping of particle motion. Drawn at random.

    Returns: next cell, indexed (i, j).
    """
    i = prev[0]   #get current cell indices  
    j = prev[1]
    
    find_cells = get_cell_action((i, j), x, y)      #compute possible moves
    K = len(np.concatenate(find_cells))        #get total num of possible moves

    rand = np.random.randint(0, K)  #generate random number between 0 and K-1 inclusive

    i_next = find_cells[0]  #separate x and y next moves 
    j_next = find_cells[1]

    if len(i_next) == len(j_next):     
        if K == 2:              #if i,j is in a corner, only need random /2
            if rand == 0:      
                i = i_next[0]          #move in x 
            if rand == 1:
                j = j_next[0]           #move in y 

        if K == 4:              #if i,j is not in corner or edge (interior), need random /4 
            if rand == 0:   
                i = i_next[0]       #move x 
            if rand == 1:
                i = i_next[1]
            if rand == 2:       #move y 
                j = j_next[0]
            if rand == 3:
                j = j_next[1]

    if K == 3:
        if len(i_next)>len(j_next):     #if j is along an edge
            if rand == 0:
                i = i_next[0]   #move x 
            if rand == 1:       
                i = i_next[1]  
            if rand == 2:          #move y 
                j = j_next[0]

        if len(i_next)<len(j_next): #if i is along an edge 
            if rand == 0:
                i = i_next[0]       #move x 
            if rand == 1:       #move y 
                j = j_next[0]
            if rand == 2:
                j = j_next[1]    

    return (i, j) #return the next propagated cell generated by random numbers




def draw_pos(history, x, y):        #write a function to draw the previous locations of the cell, weighted by how long ago it was last there 
    """
    Draw the position and history of the particle.

    Returns: (x, y) grid with current location and history of particle represented by weights out of 0 to 1.
    """
    K = len(history)        #get number of previous steps 
    N = len(x)              #grid size
    M = len(y)

    grid = np.zeros((N, M))     #generate empty grid to be populated 

    n = 0       #counter
    while n < K-1:      #repeat the iteration of drawing history until we get to the most recent position
        ip = history[n][0]      #x history 
        jp = history[n][1]      #y history 
        grid[ip][jp] = 0.6*n/K      #weight the values to be less than 1 so it's easy to follow
        n += 1                  #up counter 
            
    i = history[K-1][0]     #for the most recent cell, make it 1 to be the brightest
    j = history[K-1][1]
    grid[i][j] = 1

    return grid         #return the drawn grid 


    #generate init params 
fr = 0
cell_list = []      #keep track of cells
num = 5000          #total num of frames
i, j = (L, L)     #initial particle position
time_axis = np.linspace(0, 5, num+1)  #time axis


    #generate the data
while fr < num + 1:     #while less than frame number 
    cell_list.append((i, j))        #record cell history 
    i_next, j_next = next_cell((i, j), x, y)        #get next cell
    i = i_next
    j = j_next  #repeat 

    fr += 1 #up frame number 


if do_animations:
        #animate through the data 
    def set_axis(ax, frame):        #set the axes for matplotlib
        ax.cla()    #clear plot
        ax.set_title('Brownian Motion Sim at Time: %.2fs'%time_axis[frame], fontsize=18)
        ax.set_xlabel(r'$x$ Position (mm)', fontsize=16)
        ax.set_ylabel(r'$y$ Position (mm)', fontsize=16)

    def update(frame):      #function to update based on frame num
        set_axis(ax, frame)     #draw axes
        ax.imshow(draw_pos(cell_list[0:frame+1], x, y), extent=[0, 2*L, 0, 2*L])    #draw the most recent cell position as function of frame, moving in order from cell history
            # use the drawing function 


    fig, ax = plt.subplots(1,1)
    anim = animation.FuncAnimation(fig, update, frames = num, interval = 20)    #animate 
    # wr=animation.FFMpegWriter(fps=30)         #save the animation 
    # anim.save(r'/users/jacealloway/Desktop/brownmotion.mp4', writer=wr)
    plt.show()

        #plot x and y history 
fig, (ax1, ax2) = plt.subplots(1, 2, sharey = True)     
ax1.plot(time_axis, np.transpose(cell_list)[0]) #plot the x position vs time from the cell history 
    #labels 
ax1.set_xlabel(r'Time (s)', fontsize=16)
ax1.set_ylabel(r'Particle Path (mm)', fontsize=16) 
ax1.set_title(r'Brownian Motion $x$-Position', fontsize=18)
ax2.plot(time_axis, np.transpose(cell_list)[1]) #plot the y position vs time from the cell history 
    #labels
ax2.set_xlabel(r'Time (s)', fontsize=16)
ax2.set_title(r'Brownian Motion $y$-Position', fontsize=18)
plt.show()

        #draw 5000 iterations with previous particle locations using draw function 
plt.imshow(draw_pos(cell_list, x, y), extent=[0, 2*L, 0, 2*L])
plt.title('Brownian Motion Particle Path for %is'%int(num*0.001), fontsize=18)
plt.xlabel(r'$x$ Position (mm)', fontsize=16)
plt.ylabel(r'$y$ Position (mm)', fontsize=16)
plt.show()








"""
--------------------------------------------------- Q1 (b)
"""


    #1 for anchored location, 0 for unoccupied.
anchor = np.zeros((len(x), len(y)))


def get_cell_action(cell, x_axis, y_axis, anchor_grid):     #essentially same function as in part (a), with some anchor parameters
    """
    Returns list of possible accessible cells for an entry. Checks boundaries and anchor grid. If anchored, return True.

    Cell: array of indices.
    """
    i = cell[0]
    j = cell[1]
    N = len(x_axis)-1
    M = len(y_axis)-1


    if (i<0) or (i>=N+1) or (j<0) or (j>=M+1):
        raise ValueError('i or j must be located within grid range!')
    
    xout = []
    yout = []
    anchored = False #init


    if i == N:
        xout.append(N-1)
        anchored = True         #if along wall, anchor it 

    if i == 0:
        xout.append(i+1)
        anchored = True 

    elif 0<i<N:
        xout.append(i-1)
        xout.append(i+1)
        if (anchor_grid[i-1, j] == 1) or (anchor_grid[i+1, j] == 1):
            anchored = True         #if the anchor grid is 1 for adjacent cells, anchor it 
            

    if j == 0:
        yout.append(j+1)
        anchored = True         #anchor if along wall
    
    if j == M:
        yout.append(M-1)
        anchored = True 

    elif 0<j<M:
        yout.append(j-1)
        yout.append(j+1)
        if (anchor_grid[i, j-1] == 1) or (anchor_grid[i, j+1] == 1):
            anchored = True         #anchor if beside previously anchored particle

    return xout, yout, anchored     #return x and y next cells, and the anchored particle grid


def next_cell(prev, x, y, anchor_grid):     #almost identical to part (a) function, but now account for anchoring
    i = prev[0] 
    j = prev[1]
    
    find_cells = get_cell_action((i, j), x, y, anchor_grid)[0:2]        #get the cell action, but with the anchored particles as well 
    K = len(np.concatenate(find_cells))        #get total num of possible moves

    rand = np.random.randint(0, K)  #generate random number between 0 and K-1 inclusive

    i_next = find_cells[0]
    j_next = find_cells[1]
                                        #same as part (a)
    if len(i_next) == len(j_next):
        if K == 2:              #if i,j is in a corner
            if rand == 0:
                i = i_next[0]
            if rand == 1:
                j = j_next[0]

        if K == 4:              #if i,j is not in corner or edge 
            if rand == 0:   
                i = i_next[0]
            if rand == 1:
                i = i_next[1]
            if rand == 2:
                j = j_next[0]
            if rand == 3:
                j = j_next[1]

    if K == 3:
        if len(i_next)>len(j_next): #if j is along an edge
            if rand == 0:
                i = i_next[0]
            if rand == 1:
                i = i_next[1]
            if rand == 2:
                j = j_next[0]

        if len(i_next)<len(j_next): #if i is along an edge 
            if rand == 0:
                i = i_next[0]
            if rand == 1:
                j = j_next[0]
            if rand == 2:
                j = j_next[1]    

    return (i, j)



def draw_pos(history, x, y):        #almost same as part (a), but now accounting for all the anchored particles and drawing them 
    """
    draw the position and history of the particle.
    """
    K = len(history)
    N = len(x)
    M = len(y)

    grid = np.zeros((N, M))

    n = 0
    while n < K-1:          
        ip = history[n][0]      #get x and y pos of history entrys
        jp = history[n][1]

        if history[n][2] == 'T':          #if the particle is anchored ('T' for True), draw the grid as 1
            grid[ip][jp] = 1
        else:
            grid[ip][jp] = 0.6*n/K      #if not anchored, draw history as weights as in part (a)
        n += 1
        

    i = history[K-1][0]     #draw current most particle as 1 
    j = history[K-1][1]
    grid[i][j] = 1

    
    return np.array(grid)       #return the grid 


"""#some pseudocode
#while the center is still unanchored, repeat the loop 
#steps: 
#       1. spawn in particle at (50, 50)
#       2. propagate particle until it is anchored 
#       3. once anchored, store it in the anchored array 
#       4. repeat """

    #generate init params 
fr = 0
num = 5000          #max num of frames
time_axis = np.linspace(0, 5, num+1)  #time axis
n=0                    #set n as number of iteration steps, even when we spawn in new particle
hist = []


    #generate data 
while anchor[L][L] == 0:        #while the center position is unanchored, do this loop 
    init_cell = (L, L)              #spawn in particle in center
    current_cell = init_cell           #get the cell of the current unanchored particle
    while get_cell_action(current_cell, x, y, anchor)[2] == False:  #repeat until particle is anchored
        hist.append((current_cell[0], current_cell[1], 'F'))      #keep track of history, default unanchored is 'F' 
        i_fwd, j_fwd = next_cell(current_cell, x, y, anchor)        #propagate until anchored
        current_cell = (i_fwd, j_fwd)
        n +=1       #up the number of steps
    
    if get_cell_action(current_cell, x, y, anchor)[2] == True:  #check if the current particle is anchored; 
        i_fwd = current_cell[0]     #if it's anchored, draw it in as solid by adding it to the anchored grid
        j_fwd = current_cell[1]
        anchor[i_fwd][j_fwd] = 1                                #if so, store it in the anchor grid
    hist.append((current_cell[0], current_cell[1], 'T'))        #append the history of particles, if it's anchored, write 'T' for draw function to recognize 

    





if do_animations:
        #animate through the data 
    def set_axis(ax, frame):        #set axes
        ax.cla()    #clear plot 
        ax.set_title('Brownian Motion Anchor Sim at Time: %.2fs'%time_axis[frame], fontsize=18)
        ax.set_xlabel(r'$x$ Position (mm)', fontsize=16)
        ax.set_ylabel(r'$y$ Position (mm)', fontsize=16)

    def update(frame):
        set_axis(ax, frame)     #set axes
        ax.imshow(draw_pos(hist[0:frame+1], x, y), extent=[0, 2*L, 0, 2*L]) #use draw function with current cell, history, and anchored particles to plot everything with imshow



    fig, ax = plt.subplots(1,1)
    anim = animation.FuncAnimation(fig, update, frames = n+1, interval = 20)        #animate for number of steps + 1 (just to be safe...)
    # wr=animation.FFMpegWriter(fps=30)     #save the animation 
    # anim.save(r'/users/jacealloway/Desktop/anchors.mp4', writer=wr)
    plt.show()


    #plot the final result of all the anchored particles with the anchored grid 
plt.imshow(anchor, extent=[0, 2*L, 0, 2*L])      
    #labels
plt.title(f'Brownian Motion Anchored Particles: {n/1000}s', fontsize=18)    #the number of ms is the number of frames, convert to seconds by dividing by 1000 and use it on the plot 
plt.xlabel(r'$x$ Position (mm)', fontsize=16)
plt.ylabel(r'$y$ Position (mm)', fontsize=16)
plt.show()