import scipy as sp

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import matplotlib.pyplot as plt

import mpl_toolkits.mplot3d.axes3d as p3
import matplotlib.animation as animation

#dr  = .1
#dp  = .1
#nr      = int(1/dr)
#np      = int(2*sp.pi/dp)


nr  = 10
np  = 10

r       = sp.linspace(0.,1.,nr)
p       = sp.linspace(0.,2*sp.pi,np)

dr  = r[1]-r[0]
dp  = p[1]-p[0]

a   = .5

tmax    = 100
t       = 0.


dr2     = dr**2
dp2     = dp**2

dt      = dr2 * dp2 / (2 * a * (dr2 + dp2) )
dt      /=10.
print 'dr = ',dr 
print 'dp = ',dp 
print 'dt = ',dt



ut      = sp.zeros([nr,np])
u0      = sp.zeros([nr,np])

ur      = sp.zeros([nr,np])
ur2     = sp.zeros([nr,np])


#initial

for i in range(nr):
    for j in range(np):
        if ((i>4)&(i<6)):
            u0[i,j] = 1.
#print u0

def hitung_ut(ut,u0):
    for i in sp.arange (len(r)):
        if r[i]!= 0.:
            ur[i,:]     = u0[i,:]/r[i]
            ur2[i,:]     = u0[i,:]/(r[i]**2)
    ut[1:-1, 1:-1]  = u0[1:-1, 1:-1] + a*dt*(
            (ur[1:-1, 1:-1] - ur[:-2, 1:-1])/dr+
            (u0[2:, 1:-1] - 2*u0[1:-1, 1:-1] + u0[:-2,1:-1])/dr2+
            (ur2[1:-1, 2:] - 2*ur2[1:-1, 1:-1] + ur2[1:-1, :-2])/dp2)
    #calculate the edge
    ut[1:-1, 0]  = u0[1:-1, 0] + a*dt*(
            (ur[1:-1, 0] - ur[:-2, 0])/dr+
            (u0[2:, 0] - 2*u0[1:-1, 0] + u0[:-2, 0])/dr2+
            (ur2[1:-1, 1] - 2*ur2[1:-1, 0] + ur2[1:-1, np-1])/dp2)
    ut[1:-1, np-1]  = u0[1:-1, np-1] + a*dt*(
            (ur[1:-1, np-1] - ur[:-2, np-1])/dr+
            (u0[2:, np-1] - 2*u0[1:-1, np-1] + u0[:-2,np-1])/dr2+
            (ur2[1:-1, 0] - 2*ur2[1:-1, np-1] + ur2[1:-1, np-2])/dp2)


#hitung_ut(ut,u0)
#print ut

def data_gen(framenumber, Z ,surf):
    global ut
    global u0
    global t
    hitung_ut(ut,u0)
    u0[:] = ut[:]
    Z = u0
    t += 1
    print t
    
    ax.clear()
    plotset()
    surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False, alpha=0.7)
    return surf,


fig = plt.figure()
#ax = fig.gca(projection='3d')
ax = fig.add_subplot(111, projection='3d')

P,R = sp.meshgrid(p,r)

X,Y = R*sp.cos(P),R*sp.sin(P) 

Z = u0
print len(R), len(P)



def plotset():
    ax.set_xlim3d(-1., 1.)
    ax.set_ylim3d(-1., 1.)
    ax.set_zlim3d(-1.,1.)
    ax.set_autoscalez_on(False)
    ax.zaxis.set_major_locator(LinearLocator(10))
    ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
    cset = ax.contour(X, Y, Z, zdir='x', offset=-1. , cmap=cm.coolwarm)
    cset = ax.contour(X, Y, Z, zdir='y', offset=1. , cmap=cm.coolwarm)
    cset = ax.contour(X, Y, Z, zdir='z', offset=-1., cmap=cm.coolwarm)

plotset()
surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False, alpha=0.7)

fig.colorbar(surf, shrink=0.5, aspect=5)

ani = animation.FuncAnimation(fig, data_gen, fargs=(Z, surf),frames=4096, interval=4, blit=False)
#ani.save('2dDiffusionfRadialf1024b512.mp4', bitrate=1024)

plt.show()