import numpy as np
import sys
import math
import time
from timeit import default_timer as timer
import numpy as np
from numba import jit
from numba import cuda, float32

@cuda.jit
def vist_wale(u3d,v3d,vis3d):
   
   i, j = cuda.grid(2)

   if i < u3d.shape[0] and j <u3d.shape[1]:

      dudx=(u3d[min(i+1,u3d.shape[0]-1),j] - u3d[max(i-1,0),j])/2
      dvdx=(v3d[min(i+1,u3d.shape[0]-1),j] - v3d[max(i-1,0),j])/2

      dudy=(u3d[i,min(j+1,u3d.shape[1]-1)] - u3d[i,max(j-1,0)])/2
      dvdy=(v3d[i,min(j+1,u3d.shape[1]-1)] - v3d[i,max(j-1,0)])/2
   
      s11=dudx
      s12=0.5*(dudy+dvdx)
   
      s22=dvdy

      g11_2=dudx*dudx+dudy*dvdx
      g12_2=dudx*dudy+dudy*dvdy
   
      g21_2=dvdx*dudx+dvdy*dvdx
      g22_2=dvdx*dudy+dvdy*dvdy
   
      gkk_2=(g11_2+g22_2)/2.

      sd11=g11_2-gkk_2
      sd12=0.5*(g12_2+g21_2)
      sd22=g22_2-gkk_2
   
      sijsij=s11*s11+s12*s12+s12*s12+s22*s22

      sdijsdij=sd11*sd11+sd12*sd12+sd12*sd12+sd22*sd22

   vis3d[i,j]=sdijsdij**1.5/(sijsij**2.5+sdijsdij**1.25)

#     if i == 0 and j == 0:
#        from pdb import set_trace; set_trace()

ni=5000
nj=5000

ni=14000
nj=14000

ni=1000
nj=1000


# create velocity fields
np.random.seed(0);u3d=np.random.rand(ni,nj);v3d=np.random.rand(ni,nj)

u3d=np.float32(u3d)
v3d=np.float32(v3d)

# Copy the arrays to the device
u3d_device= cuda.to_device(u3d)
v3d_device= cuda.to_device(v3d)


# Allocate memory on the device for the result
vis3d_host= cuda.device_array((ni, nj),dtype=np.float32,)

nb_float=np.float32

##################### ni=nj=5000
threadsperblock = (16, 16)  # time = 0.35483860969543457
#threadsperblock = (32, 32) # time = 0.3948853015899658
#threadsperblock = (8, 8)  #  time = 0.4077944755554199
#threadsperblock = (4, 4)  #  time = 0.5132663249969482


##################### ni=nj=14000
threadsperblock = (32, 32) # time = 1.178175449371338
threadsperblock = (16, 16)  # time = 1.1780712604522705
threadsperblock = (8, 8) # time = 1.1637439727783203
#threadsperblock = (4, 4) # time = 1.9627246856689453



blockspergrid_x = int(math.ceil(u3d.shape[0] / threadsperblock[0]))
blockspergrid_y = int(math.ceil(u3d.shape[1] / threadsperblock[1]))
blockspergrid = (blockspergrid_x, blockspergrid_y)


start= time.time()

print('threadsperblock,blockspergrid',threadsperblock,blockspergrid)

# call vist_wale

vist_wale[blockspergrid, threadsperblock](u3d_device,v3d_device,vis3d_host)

vis3d_results = vis3d_host.copy_to_host()


print("vis3d_results[:,2]", vis3d_results[:,2])
print("\nvis3d_results[2,:]", vis3d_results[2,:])

print("time", time.time()-start)