'''
Forward method to solve 1D reaction-diffusion equation:
    u_t = D * u_xx + alpha * u

with Dirichlet boundary conditions u(x0,t) = 0, u(xL,t) = 0
and initial condition u(x,0) = 4*x - 4*x**2
'''


import numpy as np
from scipy import sparse
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
import math


M = 40 # GRID POINTS on space interval
N = 70 # GRID POINTS on time interval

x0 = 0.0
xL = 1.0

# ----- Spatial discretization step -----
dx = (xL - x0)/(M - 1)

t0 = 0.0
tF = 0.2

# ----- Time step -----
dt = (tF - t0)/(N - 1)

D = 0.1  # Diffusion coefficient
alpha = -3 # Reaction rate

r = (dt*D)/pow(dx,2)
s = dt*alpha

print("r =", r)

# ----- Creates grids -----
xspan = np.linspace(x0, xL, M)
tspan = np.linspace(t0, tF, N)

# ----- Initializes matrix solution U -----
U = np.zeros((M, N))

# ----- Initial condition -----
U[:,0] = 4*xspan - 4*xspan**2
#U[:,0] = 0.0

# ----- Dirichlet Boundary Conditions -----
U[0,:] = 0.0
U[-1,:] = 0.0 #same as U[L,:] = U[M-1,:] = 0... -1 gets the last entry

# ----- Equation (15.8) in Lecture 15 -----
for k in range(0, N-1):
    for i in range(1, M-1):
        U[i, k+1] = r*U[i-1, k] + (1-2*r+s)*U[i,k] + r*U[i+1,k]

X, T = np.meshgrid(tspan, xspan)

fig = plt.figure()
ax = fig.gca(projection='3d')

surf = ax.plot_surface(X, T, U, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False)

ax.set_xticks([0, 0.05, 0.1, 0.15, 0.2])

ax.set_xlabel('Space')
ax.set_ylabel('Time')
ax.set_zlabel('U')
plt.tight_layout()
plt.show()

import numpy as np
from scipy import sparse
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
import math

def oneDendriticBranchModel(tauV, E, lambd, i_app, i_app_location, j_input, tauR):

    M = 41 # GRID POINTS on space interval
    N = 41 # GRID POINTS on time interval

    x0 = 0.0
    xL = 20.0

    # ----- Spatial discretization step -----
    dx = (xL - x0)/(M - 1)

    t0 = 0.0
    tF = 20.0

    # ----- Time step -----
    dt = (tF - t0)/(N - 1)

    D = 1  # Diffusion coefficient
    alpha = 1 # Reaction rate

    #r = (dt*D)/pow(dx,2)
    #s = dt*alpha
    r = dt/tauV
    s = 1 - r - ((2*(pow(lambd,2))*dt)/(tauV*pow(dx,2)))
    t = (((pow(lambd,2))*dt)/(tauV*pow(dx,2)))
    indicator = 0

    print("r =", r)

    # ----- Creates grids -----
    xspan = np.linspace(x0, xL, M)
    tspan = np.linspace(t0, tF, N)

    # ----- Initializes matrix solution V -----
    V = np.zeros((M, N))

    # ----- Initial condition -----
    V[:,0] = 0

    # ----- Dirichlet Boundary Conditions -----
    V[0,:] = 0.0
    V[-1,:] = 0.0 #same as V[L,:] = V[M-1,:] = 0... -1 gets the last entry

    # ----- Equation (15.8) in Lecture 15 -----
    for j in range(0, N-1):
        for i in range(1, M-1):

            if i == i_app_location and j >= j_input and j <= j_input + tauR:
                indicator = 1
            else:
                indicator = 0

            V[i,j+1] = r*E + s*V[i,j] + t*V[i+1,j] + r*V[i-1,j] + r*i_app*indicator

    X, T = np.meshgrid(tspan, xspan)

    fig = plt.figure()
    ax = fig.gca(projection='3d')

    surf = ax.plot_surface(X, T, V, cmap=cm.coolwarm,
                           linewidth=0, antialiased=False)

    #ax.set_xticks([0, 0.05, 0.1, 0.15, 0.2])

    ax.set_xlabel('Space')
    ax.set_ylabel('Time')
    ax.set_zlabel('V')
    plt.tight_layout()
    plt.show()

oneDendriticBranchModel(1.0,0.0,1.0,2.0,1.0,1.0,2.0) #have to be float or we lose precision...