import torch


def make_complex(arr):
    new_arr = torch.stack((arr, torch.zeros(arr.shape, dtype=arr.dtype, device=arr.device)), dim=2)
    return new_arr


def tensor_to_complex(t):
    arr = t.numpy()
    return arr[...,0] + 1j*arr[...,1]


def tensor_complex_mulx(t1, t2):
    # t1, t2 should be 3d tensors where the last dim should be of size 2
    # representing the real and complex part
    arr_real = t1[...,0]*t2[...,0] -t1[...,1]*t2[...,1]
    arr_complex = t1[...,0]*t2[...,1] + t1[...,1]*t2[...,0]
    return torch.stack((arr_real, arr_complex), dim=2)


def make_conjugate(arr):
    new_arr = torch.stack((arr[:,:,0], -arr[:,:,1]), dim=2)
    return new_arr


def tensor_complex_div(t1, t2):
    denom = torch.pow(t2[:,:,0], 2) + torch.pow(t2[:,:,1], 2)
    arr_real = t1[...,0]*t2[...,0]  + t1[...,1]*t2[...,1]
    arr_real /= denom
    arr_complex = t1[..., 1] * t2[..., 0] - t1[..., 0] * t2[..., 1]
    arr_complex /=  denom
    return torch.stack((arr_real, arr_complex), dim=2)

if __name__=='__main__':
    import numpy as np
    a = np.random.random((100,100))
    b = np.random.random((100,100))
    c = np.random.random((100,100))
    d = np.random.random((100,100))
    z1 = a+1j*b
    z2 = c+1j*d
    d1 = z1/z2
    print(d1)
    z1t = torch.stack((torch.from_numpy(a), torch.from_numpy(b)), dim=2)
    z2t = torch.stack((torch.from_numpy(c), torch.from_numpy(d)), dim=2)
    d1t = tensor_complex_div(z1t, z2t)
    d1t = d1t.numpy()
    d1t = d1t[:,:,0] + 1j*d1t[:,:,1]
    print(np.allclose(d1, d1t))
