import cupy as cp
import numpy as np
import cv2


def cv2cp(mat: cv2.cuda.GpuMat) -> cp.ndarray:
    class CudaArrayInterface:
        def __init__(self, gpu_mat: cv2.cuda.GpuMat):
            w, h = gpu_mat.size()
            type_map = {
                cv2.CV_8U: "|u1",
                cv2.CV_8S: "|i1",
                cv2.CV_16U: "<u2", cv2.CV_16S: "<i2",
                cv2.CV_32S: "<i4",
                cv2.CV_32F: "<f4", cv2.CV_64F: "<f8",
                cv2.CV_16F: "<f2"
            }
            self.__cuda_array_interface__ = {
                "version": 3,
                "shape": (h, w, gpu_mat.channels()) if gpu_mat.channels() > 1 else (h, w),
                "typestr": type_map[gpu_mat.depth()],
                "descr": [("", type_map[gpu_mat.depth()])],
                "stream": 1,
                "strides": (gpu_mat.step, gpu_mat.elemSize(), gpu_mat.elemSize1()) if gpu_mat.channels() > 1
                else (gpu_mat.step, gpu_mat.elemSize()),
                "data": (gpu_mat.cudaPtr(), False),
            }
    arr = cp.asarray(CudaArrayInterface(mat))

    return arr


def cp2cv(arr: cp.ndarray) -> cv2.cuda.GpuMat:
    assert len(arr.shape) in (2, 3), "CuPy array must have 2 or 3 dimensions to be a valid GpuMat"
    type_map = {
        cp.dtype('uint8'): cv2.CV_8U,
        cp.dtype('int8'): cv2.CV_8S,
        cp.dtype('uint16'): cv2.CV_16U,
        cp.dtype('int16'): cv2.CV_16S,
        cp.dtype('int32'): cv2.CV_32S,
        cp.dtype('float32'): cv2.CV_32F,
        cp.dtype('float64'): cv2.CV_64F,
        cp.dtype('float16'): cv2.CV_16F
    }
    depth = type_map.get(arr.dtype)
    assert depth is not None, "Unsupported CuPy array dtype"
    channels = 1 if len(arr.shape) == 2 else arr.shape[2]
    mat_type = cv2.CV_MAKETYPE(depth, channels)
    mat = cv2.cuda.createGpuMatFromCudaMemory(arr.__cuda_array_interface__['shape'][1::-1],
                                              mat_type,
                                              arr.__cuda_array_interface__['data'][0])
    return mat


cp_a = cp.random.randint(1, 5, (2, 2)).astype(np.float16)
print('Original CuPy array and pointer')
print(cp_a, cp_a.__cuda_array_interface__['data'][0])
print('')
cv_a = cp2cv(cp_a)
np_a = cv_a.download()
print('')
print('GpuMat, initialized and downloaded to NumPy; and pointer')
print(np_a, cv_a.cudaPtr())
print(f'Assert GpuMat type is float16: {cv_a.type() == cv2.CV_16FC1}')
print('NumPy dtype:', np_a.dtype)
print('')
cp_a2 = cv2cp(cv_a)
print('Back to CuPy from GpuMat: value and pointer:')
print(cp_a2, cp_a2.__cuda_array_interface__['data'][0])
print('')

print('Now try to upload a float16 NumPy array to GpuMat:')
np_a = np.random.randint(1, 5, (2, 2)).astype(np.float16)
cv_a = cv2.cuda_GpuMat()
cv_a.upload(np_a)