import jax 
import kernex as kex 
import numpy as np 
import jax.numpy as jnp 
import timeit 
from functools import reduce 
import time 
from statistics import stdev,mean
import itertools
mat = lambda *args : jnp.arange(1,reduce(lambda x,y:x*y,args)+1).reshape(*args)




def test_and_time_conv2d():

  iters = 1000 

  dims = list(sorted(itertools.product([4,8,16,32,64],[16,32,64,128,256,512,1024])))
  timings = []
  
  for C,H in dims :

    print()
    
    @jax.jit
    def jax_conv2d(x,w):
      return jax.lax.conv_general_dilated(lhs = x,
                                          rhs = w,
                                          window_strides = (1,1),
                                          padding = 'SAME',
                                          dimension_numbers = ('NCHW', 'OIHW', 'NCHW'),)[0]


    @jax.jit
    @kex.kmap(kernel_size= (C,3,3),
              padding = ('valid','same','same'),
              relative = False)
    def kex_conv2d(x,w):
      return jnp.sum(x*w)  


    times = []
    x = jax.random.normal(jax.random.PRNGKey(0),(C,H,H))
    xx = x[None]
    w = jax.random.normal(jax.random.PRNGKey(0),(C,3,3))
    ww = w[None]
    np.testing.assert_allclose(kex_conv2d(x,w),jax_conv2d(xx,ww),atol=1e-3)

    # warm up
    kex_conv2d(x,w)
    jax_conv2d(xx,ww)
    
    for _ in range(iters):
      t1 = time.time()
      kex_conv2d(x,w).block_until_ready() 
      t2 = time.time()
      times +=[(t2-t1)]
    
    print(f'[benchmarking]:\tconv2d @ dim={(C,H,H)}')
    print('='*50)

    print(f'[kex] average : {mean(times)*1e6:.0f} us\t stddev : {stdev(times)*10e6:.3f}us')

    times = []
    for _ in range(iters):
      t1 = time.time()
      jax_conv2d(xx,ww).block_until_ready() 
      t2 = time.time()
      times +=[t2-t1]

    
    
    print(f'[jax] average : {mean(times)*1e6:.0f} us\t stddev : {stdev(times)*10e6:.3f}us')

    print('='*50)

