import numpy

  test_tensor = tf.constant([[[0,0,255]] * 1024] * 512)

  test_tensor = tf.math.scalar_mul(255, tf.keras.activations.sigmoid(inputs_equirectangular(1024, 512)))

def inputs(width, height):
    def cartesian(x,y):
        x = x / width * math.pi
        y = y / height * math.pi
        return [math.sin(y) * math.cos(x), math.sin(y) * math.sin(x), math.cos(y)]
    def row(y):
        def c(x):
            return cartesian(x,y)
        return list(map(c, range(0, width)))
    return tf.constant(list(map(row, range(0, height))))

def inputs_equirectangular(width, height):
    n = numpy.zeros((height, width, 3), numpy.float64)
    helpers.equirectangular(ctypes.cast(ctypes.c_voidp(n.ctypes.data), ctypes.POINTER(ctypes.c_double)), ctypes.c_int(width), ctypes.c_int(height))
    return tf.constant(n)

native.equirectangular.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_int, ctypes.c_int]
equirectangular = native.equirectangular

#[no_mangle]
pub extern "C" fn equirectangular(target: *mut f64, width: u32, height: u32) {
    use std::f64::consts::PI;
    let target = unsafe {
        std::slice::from_raw_parts_mut(target, (width as usize) * (height as usize) * 3)
    };
    for y in 0 .. (height as usize) {
        let rowstart = y * (width as usize) * 3;
        let yr = 0. - (y as f64) / (height as f64) * PI;
        let yrs = yr.sin();
        let yrc = yr.cos();
        for x in 0 .. (width as usize) {
            let xr = (x as f64) / (width as f64) * 2. * PI;
            let pixstart = rowstart + 3 * x;
            target[pixstart] = yrs * xr.sin();
            target[pixstart + 1] = yrs * xr.cos();
            target[pixstart + 2] = yrc;
        }
    }
}