import matplotlib.pyplot as plt
import numpy as np
import time
from numbapro import autojit, cuda, jit, float

###1-D Nonlinear convection implemented using Numpy
def NonLinNumpy(u, un, nx, nt, dx, dt):

    ###Run through nt timesteps and plot/animate each step
    for n in range(nt): ##loop across number of time steps
        un[:] = u[:]
        u[1:] = -un[1:]*dt/dx*(un[1:]-un[:-1])+un[1:]
    
    return u

###1-D Nonlinear convection implemented using 'vanilla' Python
def NonLinVanilla(u, nx, nt, dx, dt):

    for n in range(nt):
        for i in range(1,nx-1):
            u[i+1] = -u[i]*dt/dx*(u[i]-u[i-1])+u[i]

    return u

###1-D Nonlinear convection implemented using Numba JIT compiler (similar to LLVM)
@autojit
def NonLinNumba(u,un, nx, nt, dx, dt):

    for n in range(nt):
        for i in range(1,nx):
            un[i] = -u[i]*dt/dx*(u[i]-u[i-1])+u[i]

    return un

###1-D Nonlinear convection implemented using NumbaPro CUDA-JIT
d@jit(argtypes=[float32[:], float32, float32, float32, float32[:]], target='gpu')
def NonLinCudaJit(u, dx, dt, nt, un):
    tid = cuda.threadIdx.x
    blkid = cuda.blockIdx.x
    blkdim = cuda.blockDim.x
    i = tid + blkid * blkdim

    if i >= u.shape[0]:
        return

    for n in range(nt):
        un[i] = -u[i]*dt/dx*(u[i]-u[i-1])+u[i]
        
        cuda.syncthreads()

def main(nx):
    ##System Conditions    
    #nx = 500 
    nt = 500
    c = 1
    xmax = 15.0
    dx = xmax/(nx-1)
    sigma = 0.25
    dt = sigma*dx

    ##Initial Conditions for wave
    ui = np.ones(nx) ##create a 1xn vector of 1's
    ui[.5/dx:1/dx+1]=2 ##set hat function I.C. : .5<=x<=1 is 2
    un = np.ones(nx)    

    if nx < 20001:
        t1 = time.time()
        u = NonLinVanilla(ui, nx, nt, dx, dt)
        t2 = time.time()
        print "Vanilla version took: %.6f seconds" % (t2-t1)
    
    
    ui = np.ones(nx) ##create a 1xn vector of 1's
    ui[.5/dx:1/dx+1]=2 ##set hat function I.C. : .5<=x<=1 is 2
    
    t1 = time.time()
    u = NonLinNumpy(ui, un, nx, nt, dx, dt)
    t2 = time.time()
    print "Numpy version took: %.6f seconds" % (t2-t1)
    numpytime = t2-t1
    #plt.plot(numpy.linspace(0,xmax,nx),u[:],marker='o',lw=2)

    
    ui = np.ones(nx) ##create a 1xn vector of 1's
    ui[.5/dx:1/dx+1]=2 ##set hat function I.C. : .5<=x<=1 is 2
    
    t1 = time.time()
    u = NonLinNumba(ui, un, nx, nt, dx, dt)
    t2 = time.time()
    print "Numbapro Vectorize version took: %.6f seconds" % (t2-t1)
    vectime = t2-t1
    #plt.plot(numpy.linspace(0,xmax,nx),u[:],marker='o',lw=2)

    u = np.ones(nx)
    u[:] = ui[:]
    griddim = 320, 1
    blockdim = 768, 1, 1
    NonLinCudaJit_conf = NonLinCudaJit[griddim, blockdim]
    t1 = time.time()
    NonLinCudaJit(u, dx, dt, nt, un)
    t2 = time.time()

    print "Numbapro Cuda version took: %.6f seconds" % (t2-t1)
    cudatime = t2-t1

main(500)

main(1000)

main(5000)

main(10000)

main(20000)

main(50000)

main(100000)

main(200000)

from IPython.core.display import HTML
def css_styling():
    styles = open("../styles/custom.css", "r").read()
    return HTML(styles)
css_styling()