import torch
import numpy as np
import cupy as cp
from collections import namedtuple
from timeit import timeit

def timer(fn, number=1000):
    print(timeit(fn, number=number) / number)

Stream = namedtuple('Stream', ['ptr'])

#%%##################################################################
kernel = cp.RawKernel(r'''
extern "C" __global__
void cdiv_imag(float x[240][240000][2], float y[240][240000][2],
               float out[240][240000]) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;
    if (i >= 240 || j >= 240000)
        return;

    float A = x[i][j][0];
    float B = x[i][j][1];
    float C = y[i][j][0];
    float D = y[i][j][1];

    out[i][j] = (B*C - A*D) / (C*C + D*D);
}

''', 'cdiv_imag')
#%%##################################################################
shape = (240, 240000)
x = torch.randn(shape, dtype=torch.complex64, device='cuda')
y = torch.fliplr(x.clone().detach())

threadsperblock = (32, 32)
blockspergrid_x = int(np.ceil(y.shape[0] / threadsperblock[0]))
blockspergrid_y = int(np.ceil(y.shape[1] / threadsperblock[1]))
blockspergrid = (blockspergrid_x, blockspergrid_y)

#%%##################################################################
def f1(x, y):
    out = torch.zeros(shape, dtype=torch.float32, device='cuda')
    xv, yv = torch.view_as_real(x), torch.view_as_real(y)
    kernel(
        grid=blockspergrid, block=threadsperblock,
        args=[xv.data_ptr(), yv.data_ptr(), out.data_ptr()],
        stream=Stream(ptr=torch.cuda.current_stream().cuda_stream))
    return out

def f2(x, y):
    return (x / y).imag

# ensure outputs agree within float precision
adiff = torch.abs(f1(x, y) - f2(x, y))
assert adiff.mean() < 1e-8

# dummy run to factor out caching & warm up GPU
for _ in range(2000):
  f1(x, y)
  f2(x, y)

#%%##################################################################
timer(lambda: f1(x, y))
timer(lambda: f2(x, y))