import numpy
from matplotlib import pyplot

%matplotlib inline

J = 5
I = 10

print numpy.array([[str((j,i)) for j in range(J)] for i in range(I-1,-1,-1)])

L_x = 1.
L_y = 2.
T = 1.

J = 100
I = 200
N = 1000

dx = float(L_x)/float(J-1)
dy = float(L_y)/float(I-1)
dt = float(T)/float(N-1)

x_grid = numpy.array([j*dx for j in range(J)])
y_grid = numpy.array([i*dy for i in range(I)])
t_grid = numpy.array([n*dt for n in range(N)])

D_U = 0.1
D_V = 10.

tot_protein = 2.26

no_high = 10
U =  numpy.array([[0.1 if (I-1)-i+1 > no_high else 2.0 for j in range(J)] for i in range(I)])

U_protein = sum(sum(U))*dx*dy
V_protein_px = float(tot_protein-U_protein)/float(I*J*dx*dy)
V = numpy.array([[V_protein_px for i in range(0,J)] for i in range(I)])

print 'Initial protein mass', (sum(sum(U))+sum(sum(V)))*(dx*dy)

fig, ax = pyplot.subplots()
ax.set_xlim(left=0., right=(J-1)*dx)
ax.set_ylim(bottom=0., top=(I-1)*dy)
heatmap = ax.pcolor(x_grid, y_grid, U, vmin=0., vmax=2.1)
colorbar = pyplot.colorbar(heatmap)
colorbar.set_label('concentration U')

alpha_x_U = D_U*dt/(2.*dx*dx)
alpha_y_U = D_U*dt/(2.*dy*dy)

alpha_x_V = D_V*dt/(2.*dx*dx)
alpha_y_V = D_V*dt/(2.*dy*dy)

A_U = numpy.diagflat([-alpha_x_U for j in range(J-1)], -1)+\
    numpy.diagflat([1.+alpha_x_U]+[1.+2.*alpha_x_U for j in range(J-2)]+[1.+alpha_x_U], 0)+\
    numpy.diagflat([-alpha_x_U for j in range(J-1)], 1)
    
C_U = numpy.diagflat([-alpha_y_U for i in range(I-1)], -1)+\
      numpy.diagflat([1.+alpha_y_U]+[1.+2.*alpha_y_U for i in range(I-2)]+[1.+alpha_y_U], 0)+\
      numpy.diagflat([-alpha_y_U for i in range(I-1)], 1)
        
A_V = numpy.diagflat([-alpha_x_V for j in range(J-1)], -1)+\
    numpy.diagflat([1.+alpha_x_V]+[1.+2.*alpha_x_V for j in range(J-2)]+[1.+alpha_x_V], 0)+\
    numpy.diagflat([-alpha_x_V for j in range(J-1)], 1)
    
C_V = numpy.diagflat([-alpha_y_V for i in range(I-1)], -1)+\
      numpy.diagflat([1.+alpha_y_V]+[1.+2.*alpha_y_V for i in range(I-2)]+[1.+alpha_y_V], 0)+\
      numpy.diagflat([-alpha_y_V for i in range(I-1)], 1)

f_vec = lambda U, V: numpy.multiply(float(dt)/float(2.), numpy.subtract(numpy.multiply(V, 
                     numpy.add(k0, numpy.divide(numpy.multiply(U,U), numpy.add(1., numpy.multiply(U,U))))), U))

k0 = 0.067

b_t_stencil_U = numpy.array([[(1.-alpha_y_U) for j in range(J)],
                             [alpha_y_U for j in range(J)]])
b_c_stencil_U = numpy.array([[alpha_y_U for j in range(J)],
                             [1.-2.*alpha_y_U for j in range(J)],
                             [alpha_y_U for j in range(J)]])
b_b_stencil_U = numpy.array([[alpha_y_U for j in range(J)],
                             [(1.-alpha_y_U) for j in range(J)]])

b_t_stencil_V = numpy.array([[(1.-alpha_y_V) for j in range(J)],
                             [alpha_y_V for j in range(J)]])
b_c_stencil_V = numpy.array([[alpha_y_V for j in range(J)],
                             [1.-2.*alpha_y_V for j in range(J)],
                             [alpha_y_V for j in range(J)]])
b_b_stencil_V = numpy.array([[alpha_y_V for j in range(J)],
                             [(1.-alpha_y_V) for j in range(J)]])

d_l_stencil_U = numpy.array([[1.-alpha_x_U, alpha_x_U] for i in range(I)])
d_c_stencil_U = numpy.array([[alpha_x_U, 1.-2.*alpha_x_U, alpha_x_U] for i in range(I)])
d_r_stencil_U = numpy.array([[alpha_x_U,1.-alpha_x_U] for i in range(I)])
                             
d_l_stencil_V = numpy.array([[1.-alpha_x_V, alpha_x_V] for i in range(I)])
d_c_stencil_V = numpy.array([[alpha_x_V, 1.-2.*alpha_x_V, alpha_x_V] for i in range(I)])
d_r_stencil_V = numpy.array([[alpha_x_V, 1.-alpha_x_V] for i in range(I)])

f_curr = f_vec(U,V)

def b_U(i):
    if i <= I-2 and i >= 1:
        U_y = U[[i+1, i, i-1], :]
        return numpy.add(U_y+b_c_stencil_U, f_curr[[i+1, i, i-1], :]).sum(axis=0)
    elif i == I-1:
        U_y = U[[I-1, I-2], :]
        return numpy.add(U_y+b_t_stencil_U, f_curr[[I-1, I-2], :]).sum(axis=0)
    elif i == 0:
        U_y = U[[1, 0], :]
        return numpy.add(U_y+b_b_stencil_U, f_curr[[1, 0], :]).sum(axis=0)
    
def b_V(i):
    if i <= I-2 and i >= 1:
        V_y = V[[i+1, i, i-1], :]
        return numpy.subtract(V_y+b_c_stencil_V, f_curr[[i+1, i, i-1], :]).sum(axis=0)
    elif i == I-1:
        V_y = V[[I-1, I-2], :]
        return numpy.subtract(V_y+b_t_stencil_V, f_curr[[I-1, I-2], :]).sum(axis=0)
    elif i == 0:
        V_y = V[[1, 0], :]
        return numpy.subtract(V_y+b_b_stencil_V, f_curr[[1, 0], :]).sum(axis=0)
    
def d_U(j):
    if j <= J-2 and j >= 1:
        U_x = U[:, [j-1, j, j+1]]
        return numpy.add(U_x+d_c_stencil_U, f_curr[:, [j-1, j, j+1]]).sum(axis=1)
    if j == 0:
        U_x = U[:, [0, 1]]
        return numpy.add(U_x+d_l_stencil_U, f_curr[:, [0, 1]]).sum(axis=1)
    if j == J-1:
        U_x = U[:, [J-2, J-1]]
        return numpy.add(U_x+d_r_stencil_U, f_curr[:, [J-2, J-1]]).sum(axis=1)
    
def d_V(j):
    if j <= J-2 and j >= 1:
        V_x = V[:, [j-1, j, j+1]]
        return numpy.subtract(V_x+d_c_stencil_V, f_curr[:, [j-1, j, j+1]]).sum(axis=1)
    if j == 0:
        V_x = V[:, [0, 1]]
        return numpy.subtract(V_x+d_l_stencil_V, f_curr[:, [0, 1]]).sum(axis=1)
    if j == J-1:
        V_x = V[:, [J-2, J-1]]
        return numpy.subtract(V_x+d_r_stencil_V, f_curr[:, [J-2, J-1]]).sum(axis=1)

for n in range(N):
    print 'time = %.6f' % (n*dt)
    print 'total protein = %.6f' % ((dx*dy)*(sum(sum(U))+sum(sum(V))))
    f_curr = f_vec(U,V)
    
    for i in range(I):
        U[i, :] = numpy.linalg.solve(A_U, b_U(i))
        V[i, :] = numpy.linalg.solve(A_V, b_V(i))
    for j in range(J):
        U[:, j] = numpy.linalg.solve(C_U, d_U(j))
        V[:, j] = numpy.linalg.solve(C_V, d_V(j))


fig, ax = pyplot.subplots()
ax.set_xlim(left=0., right=(J-1)*dx)
ax.set_ylim(bottom=0., top=(I-1)*dy)
heatmap = ax.pcolor(x_grid, y_grid, U, vmin=0., vmax=2.1)
colorbar = pyplot.colorbar(heatmap)
colorbar.set_label('concentration U')