#include "cl_viv_vx_ext.h"

_viv_uniform int bOrder;
_viv_uniform int rOrder;

_viv_uniform float outputScaleVar;
_viv_uniform float bMeanScaleVarZp;
_viv_uniform float gMeanScaleVarZp;
_viv_uniform float rMeanScaleVarZp;

_viv_uniform uint  xrIntFloat_16;
_viv_uniform uint  yrIntFloat_16;

_viv_uniform VXC_512Bits uniConvertNV12toB_4x4;
_viv_uniform VXC_512Bits uniConvertNV12toG_4x4;
_viv_uniform VXC_512Bits uniConvertNV12toR_4x4;

_viv_uniform VXC_512Bits uniConvertHalftoFp16_2x8;

_viv_uniform VXC_512Bits uniConvertInt32toUint8_2x8;
_viv_uniform VXC_512Bits uniConvertUVtoCharSub128_2x8;

_viv_uniform VXC_512Bits uniCalculateYShift_2x8;
_viv_uniform VXC_512Bits uniCalculateUVShift_2x8;

__kernel void pre_process_nv12_scale_U8toU8_gq(
    __read_only image2d_t y_img, __read_only image2d_t uv_img,
    __write_only image2d_array_t    output,
    global int *xRatio, global int * yRatio, global int * xOffset, global int * yOffset,
    float rMean, float gMean, float bMean, float var, int reverse_channel, int trans)
{
    uint4 gidx = get_global_id(0);
    uint gidy = get_global_id(1);
    gidx += (uint4)(0, 1, 2, 3);

    uint dy = (gidy * yrIntFloat_16) >> 16;
    uint4 dx = (gidx * xrIntFloat_16) >> 16;
    int sy = convert_int(dy) + (*yOffset);
    int4 sx = convert_int4(dx) + (*xOffset);
    int4 uvX = sx & 0xfffffffe;
    int uvY = sy >> 1;

    vxc_uchar16 Y, UV;
    int2 coord = (int2)(sx.x, sy);
    int2 coord_uv = (int2)(uvX.x, uvY);
    VXC_ReadImage(Y, y_img, coord, VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
    VXC_ReadImage(UV, uv_img,coord_uv, VXC_5BITOFFSET_XY(0,0),VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));

    vxc_uchar16 maskShift = {8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8};
    vxc_uchar16 maskShiftUv = {8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8};
    int4 offsetUV = uvX - uvX.x;

    vxc_ushort8 diffY, diffUV;
    _viv_asm(COPY, diffY, sx, 16);
    _viv_asm(COPY, diffUV, offsetUV, 16);

    vxc_ushort8 constData = 8;
    VXC_DP2x8(maskShift, diffY, constData, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 1), uniCalculateYShift_2x8);
    VXC_DP2x8(maskShiftUv, diffUV, constData, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniCalculateUVShift_2x8);
    VXC_BitExtract(Y, Y, Y, maskShift, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
    VXC_BitExtract(UV, UV, UV, maskShiftUv, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    vxc_char16 tmpUV;
    short tmpVal = 128;
    VXC_DP2x8(tmpUV, UV, tmpVal, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertUVtoCharSub128_2x8);

    float4 tmpDstB, tmpDstG, tmpDstR;
    VXC_DP4x4(tmpDstB, Y, tmpUV, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertNV12toB_4x4);
    VXC_DP4x4(tmpDstG, Y, tmpUV, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertNV12toG_4x4);
    VXC_DP4x4(tmpDstR, Y, tmpUV, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertNV12toR_4x4);

    int4 result;
    vxc_uchar8 dst;
    int4 dstPos = (int4)(get_global_id(0), gidy, 0, 0);
    result = convert_int4_rte(tmpDstB * outputScaleVar + bMeanScaleVarZp);
    dstPos.z = bOrder;
    VXC_DP2x8(dst, result, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniConvertInt32toUint8_2x8);
    VXC_WriteImage2DArray(output, dstPos, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));

    result = convert_int4_rte(tmpDstG * outputScaleVar + gMeanScaleVarZp);
    dstPos.z = 1;
    VXC_DP2x8(dst, result, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniConvertInt32toUint8_2x8);
    VXC_WriteImage2DArray(output, dstPos, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));

    result = convert_int4_rte(tmpDstR * outputScaleVar + rMeanScaleVarZp);
    dstPos.z = rOrder;
    VXC_DP2x8(dst, result, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniConvertInt32toUint8_2x8);
    VXC_WriteImage2DArray(output, dstPos, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
}

__kernel void pre_process_nv12_scale_U8toF16_gq(
    __read_only image2d_t y_img, __read_only image2d_t uv_img,
    __write_only image2d_array_t    output,
    global int *xRatio, global int * yRatio, global int * xOffset, global int * yOffset,
    float rMean, float gMean, float bMean, float var, int reverse_channel, int trans)
{
    uint4 gidx = get_global_id(0);
    uint gidy = get_global_id(1);
    gidx += (uint4)(0, 1, 2, 3);

    uint dy = (gidy * yrIntFloat_16) >> 16;
    uint4 dx = (gidx * xrIntFloat_16) >> 16;
    int sy = convert_int(dy) + (*yOffset);
    int4 sx = convert_int4(dx) + (*xOffset);
    int4 uvX = sx & 0xfffffffe;
    int uvY = sy >> 1;

    vxc_uchar16 Y, UV;
    int2 coord = (int2)(sx.x, sy);
    int2 coord_uv = (int2)(uvX.x, uvY);
    VXC_ReadImage(Y, y_img, coord, VXC_5BITOFFSET_XY(0, 0), VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));
    VXC_ReadImage(UV, uv_img,coord_uv, VXC_5BITOFFSET_XY(0,0),VXC_MODIFIER(0, 15, 0, VXC_RM_TowardZero, 0));

    vxc_uchar16 maskShift = {8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8};
    vxc_uchar16 maskShiftUv = {8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8};
    int4 offsetUV = uvX - uvX.x;

    vxc_ushort8 diffY, diffUV;
    _viv_asm(COPY, diffY, sx, 16);
    _viv_asm(COPY, diffUV, offsetUV, 16);

    vxc_ushort8 constData = 8;
    VXC_DP2x8(maskShift, diffY, constData, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 1), uniCalculateYShift_2x8);
    VXC_DP2x8(maskShiftUv, diffUV, constData, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 1), uniCalculateUVShift_2x8);
    VXC_BitExtract(Y, Y, Y, maskShift, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
    VXC_BitExtract(UV, UV, UV, maskShiftUv, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));

    vxc_char16 tmpUV;
    short tmpVal = 128;
    VXC_DP2x8(tmpUV, UV, tmpVal, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertUVtoCharSub128_2x8);

    float4 tmpDstB, tmpDstG, tmpDstR;
    VXC_DP4x4(tmpDstB, Y, tmpUV, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertNV12toB_4x4);
    VXC_DP4x4(tmpDstG, Y, tmpUV, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertNV12toG_4x4);
    VXC_DP4x4(tmpDstR, Y, tmpUV, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniConvertNV12toR_4x4);

    tmpDstB = tmpDstB * outputScaleVar + bMeanScaleVarZp;
    tmpDstG = tmpDstG * outputScaleVar + gMeanScaleVarZp;
    tmpDstR = tmpDstR * outputScaleVar + rMeanScaleVarZp;

    half4 result;
    vxc_half8 tmpdst;
    vxc_short8 dst;
    int4 dstPos = (int4)(get_global_id(0), gidy, 0, 0);
    _viv_asm(CONV, result, tmpDstB);
    dstPos.z = bOrder;
    VXC_DP2x8(tmpdst, result, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8);
    _viv_asm(COPY, dst, tmpdst, 16);
    VXC_WriteImage2DArray(output, dstPos, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));

    _viv_asm(CONV, result, tmpDstG);
    dstPos.z = 1;
    VXC_DP2x8(tmpdst, result, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8);
    _viv_asm(COPY, dst, tmpdst, 16);
    VXC_WriteImage2DArray(output, dstPos, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));

    _viv_asm(CONV, result, tmpDstR);
    dstPos.z = rOrder;
    VXC_DP2x8(tmpdst, result, result, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0), uniConvertHalftoFp16_2x8);
    _viv_asm(COPY, dst, tmpdst, 16);
    VXC_WriteImage2DArray(output, dstPos, dst, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0));
}