"""
yfi: y fire index
yfs: y fire shape
xfi: x fire index
xfs: x fire shape
xri: x receive index
xrs: x receive shape
"""
import cupy as cp
from flandre.beamformer.dist import direct_dist

from flandre.utils.Config import DeviceConfig


def dist_mat_to_yids(dist_mat: cp.ndarray, dev_cfg=DeviceConfig()):
    def yids(offset: int) -> cp.ndarray:
        yfs = dev_cfg.rows
        xfs = 64
        xss = dev_cfg.cols
        xrs = dev_cfg.cols
        kernel = cp.RawKernel(f"""
                extern "C" __global__ void k(unsigned short int *result, float * dist_mat, int offset)
                {{
                    const long long cols = {dev_cfg.cols};
                    
                    const long long yfi = blockIdx.x;
                    const long long yfs = gridDim.x;
                    
                    const long long xfi = blockIdx.y;
                    const long long xfs = gridDim.y;
                    
                    const long long xsi = blockIdx.z;
                    const long long xss = gridDim.z;
                    
                    const long long xri = threadIdx.x;
                    const long long xrs = blockDim.x;

                    long long idx = xri + xsi*xrs + xfi*xrs*xss + yfi*xrs*xss*xfs;
                    float r = dist_mat[yfi*cols+abs(xsi-(xfi + offset))] + dist_mat[yfi*cols+abs(xri-(xfi + offset))];
                    result[idx] = min(__float2int_rd(r), (int)yfs);
                }}
        """, 'k')
        result = cp.zeros(yfs * xfs * xss * xrs, dtype=cp.uint16)
        kernel((yfs, xfs, xss), (xrs, 1, 1), (
            result,
            cp.asarray(dist_mat.astype(cp.float32)),
            cp.int32(offset)
        ))
        result = result.reshape((yfs, xfs, xss, xrs))
        return result

    return yids


def dist_mat_to_yids_pwi(dist_mat: cp.ndarray, dev_cfg=DeviceConfig()):
    """
    min in cuda is for __/\\__
    :param dist_mat: cp.ndarray from dist.py
    :param dev_cfg: DeviceConfig
    :return: pixel distance (yf,xf,xr)
    """

    def yids() -> cp.ndarray:
        yfs = dev_cfg.rows
        xfs = dev_cfg.cols
        xrs = dev_cfg.cols
        kernel = cp.RawKernel(f"""
                extern "C" __global__ void k(unsigned short int *result, float * dist_mat)
                {{
                    const long long cols = {dev_cfg.cols};
                    
                    const long long yfi = blockIdx.x;
                    const long long yfs = gridDim.x;
                    
                    const long long xfi = blockIdx.y;
                    const long long xfs = gridDim.y;
                    
                    const long long xri = threadIdx.x;
                    const long long xrs = blockDim.x;

                    long long idx = xri + xfi*xrs + yfi*xrs*xfs; //(yf,xf,xr)
                    float r = dist_mat[yfi*cols] + dist_mat[yfi*cols+abs(xri-xfi)];
                    result[idx] = min(__float2int_rd(r), (int)yfs); 
                }}
        """, 'k')
        result = cp.zeros(yfs * xfs * xrs, dtype=cp.uint16)
        kernel((yfs, xfs), (xrs, 1, 1), (result, cp.asarray(dist_mat.astype(cp.float32))))
        result = result.reshape((yfs, xfs, xrs))
        return result

    return yids


def test1():
    f = dist_mat_to_yids(direct_dist())
    r = f(0)
    print(r.flags['C_CONTIGUOUS'], r.flags['F_CONTIGUOUS'])


if __name__ == '__main__':
    test1()
