import numpy as np

nbframes = 50 # saving "nbframes" frames

dx = 0.001 # in mm
dt = 0.02 # in minutes

L = 9 # size of spatial domain, i.e. [0,L]
T = 1500 # time of simulation

b0 = 0.01 # oxygen consumption rate per cell
a0 = 0.001 # aerotactic advection speed
Drho = 3/(10**5) # diffusivity of cells
Doxy = 0.12 # diffusivity of oxygen
r = np.log(2)/440 # cell division rate
Ninit = 21 # initial oxygen concentration
Nt2 = 0.7 # first oxygen threshold C_0
Nt1 = 0.1 # first oxygen threshold C_0

N = np.load("N_initial.npy") # oxygen field
rho = np.load("rho_initial.npy") # cell density field

nbsteps = int(np.floor(T/dt))
I = int(np.floor(L/dx))
Irho = int(I/3) # rho is computed on the domain [0,3mm], in order to speed up computations
everynthframe = int(np.ceil(nbsteps/nbframes))
countframe = 0

N_data = np.zeros((nbframes,I)) # saves oxygen field after each everynthframe iterations
rho_data = np.zeros((nbframes,I)) # saves cell density field after each everynthframe iterations
neutralfraction_data = np.zeros((nbframes,Irho)) # saves cell density of the labeled cells afer each everynthframe iterations
 
neutralfraction = np.copy(rho) # cell density of the labeled cells
neutralfraction[270:] *= 0 # only cells below i=270 get labeled

def chi_array(I, N, Nt2, a0): # returns the field a(x) as an array
    res = np.zeros(I)
    for i in range(I):
        Naux = N[i]
        if Naux<=Nt2:
            res[i] = a0
    return(res)

def r_array(I, N, Nt2, r): # returns the field r(x) as an array
    res = np.zeros(I)
    for i in range(I):
        Naux = N[i]
        if Nt2<=Naux:
            res[i] = r
    return(res)

def thomas_tridiag(a, b, c, d, I): #Solves the system Mx=d where M=tridiag(a,b,c)
    res = np.zeros(I)
    cp = np.zeros(I)
    dp = np.zeros(I)
    cp[0] = c[0]/b[0]
    for i in range(1,I-1):
        cp[i] = c[i]/(b[i]-a[i]*cp[i-1])
    dp[0] = d[0]/b[0]
    for i in range(1,I):
        dp[i] = (d[i]-a[i]*dp[i-1])/(b[i]-a[i]*cp[i-1])
    res[I-1] = dp[I-1]
    for i in range(I-2,-1,-1):
        res[i] = dp[i]-cp[i]*res[i+1]
    return(res)

for t in range(nbsteps):
    # Saving data
    if t%everynthframe==0: 
        time = countframe*everynthframe*dt
        N_data[countframe,:] = N
        rho_data[countframe,:] = rho
        neutralfraction_data[countframe,:] = neutralfraction
        countframe += 1
    # Scheme for rho
    chi_field = chi_array(Irho, N, Nt2, a0)
    r_field = r_array(Irho, N, Nt2, r)
    a = np.zeros(Irho)
    b = np.zeros(Irho)
    c = np.zeros(Irho)
    i = 0
    a[i] = 0
    b[i] = 1+Drho*dt/(dx*dx)-dt*r_field[i] # homogeneous Neumann BC
    c[i] = -Drho*dt/(dx*dx)
    for i in range(1,Irho-1):
        a[i] = -Drho*dt/(dx*dx)-dt/dx*chi_field[i-1]
        b[i] = 1+2*Drho*dt/(dx*dx)-dt*r_field[i]+dt/dx*chi_field[i]
        c[i] = -Drho*dt/(dx*dx)
    i = Irho-1
    a[i] = -Drho*dt/(dx*dx)-dt/dx*chi_field[i-1]
    b[i] = 1+Drho*dt/(dx*dx)-dt*r_field[i]+dt/dx*chi_field[i] # homogeneous Neumann BC
    c[i] = 0
    rho = thomas_tridiag(a, b, c, rho, Irho)
    neutralfraction = thomas_tridiag(a, b, c, neutralfraction, Irho)
    # scheme for N
    a = np.zeros(I)
    b = np.zeros(I)
    c = np.zeros(I)
    e = np.zeros(I)
    a = -Doxy*dt/(dx*dx)*np.ones(I)
    b[0] = 1+Doxy*dt/(dx*dx)+dt*b0*rho[0]*(N[0]<Nt1)*N[0]/Nt1 # homogeneous Neumann BC
    b[1:Irho] = (1+2*Doxy*dt/(dx*dx))*np.ones(Irho-1)+dt*b0*rho[1:]*(N[1:Irho]<Nt1)*N[1:Irho]/Nt1
    b[Irho:] = (1+2*Doxy*dt/(dx*dx))*np.ones(I-Irho)
    c = -dt*Doxy/(dx*dx)*np.ones(I)
    e[:Irho] += -dt*b0*rho*(N[:Irho]>Nt1)
    e[-1] += Doxy*dt/(dx*dx)*Ninit # Dirichlet BC N=Ninit
    N = thomas_tridiag(a, b, c, N+e, I)

# Saving data

# Loading data.npz
data_dic = dict(np.load('../../data.npz'))
# Replacing data
data_dic['N_data_AC'] = N_data
data_dic['rho_data_AC'] = rho_data
data_dic['neutralfraction_data_AC'] = neutralfraction
# Saving present data.npz
np.savez('../../data.npz',**data_dic)

    
    

    
    
    
