#include "cl_viv_vx_ext.h"
_viv_uniform int class_max_iter;

_viv_uniform VXC_512Bits uniPackMaxData_2x8;
_viv_uniform VXC_512Bits uniGetSubData0to3_4x4;
_viv_uniform VXC_512Bits uniGetSubData4to7_4x4;
_viv_uniform int iter;
_viv_uniform int stride;
_viv_uniform float re_rand_max;

uint4 _philox4x32bumpkey(uint4 key)
{
    uint4 mask = (uint4)((uint)0x9E3779B9, (uint)0xBB67AE85, 0, 0);
    //key.x += ((uint)0x9E3779B9);
    //key.y += ((uint)0xBB67AE85);
    key += mask;
    return key;
}

uint mulhilo32(uint a, uint b, uint* hip)
{
    uint product = (uint)(a * b);
    *hip = mul_hi(a, b);
    return product;
}

uint mullo32(uint a, uint b)
{
    return a * b;
}

uint mulhi32(uint a, uint b)
{
    return mul_hi(a, b);
}

uint4 _philox4x32round(uint4 ctr, uint4 key)
{
    //uint hi0;
    //uint hi1;
    uint PHILOX_M4x32_0 = ((uint)0xD2511F53);
    uint PHILOX_M4x32_1 = ((uint)0xCD9E8D57);
    uint lo0 = mullo32(PHILOX_M4x32_0, ctr.x);
    uint hi0 = mulhi32(PHILOX_M4x32_0, ctr.x);
    uint lo1 = mullo32(PHILOX_M4x32_1, ctr.z);
    uint hi1 = mulhi32(PHILOX_M4x32_1, ctr.z);
    //uint lo0 = mulhilo32(PHILOX_M4x32_0, ctr.x, &hi0);
    //uint lo1 = mulhilo32(PHILOX_M4x32_1, ctr.z, &hi1);
    uint4 out = (uint4)(hi1^ctr.y^key.x, lo1, hi0^ctr.w^key.y, lo0);
    return out;
}

uint4 philox4x32_R_10(uint4 ctr, uint4 key)
{
    uint i;
    ctr = _philox4x32round(ctr, key);
    for (i = 1; i < 10; i++)
    {
        key = _philox4x32bumpkey(key);
        ctr = _philox4x32round(ctr, key);
    }
    return ctr;
}

__kernel void random_seed(
    __read_only  image2d_t  seeds,
    __write_only image2d_t  output)
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int4 coord = (int4)(gidx << 1, gidy, 0, 0);

    int width = get_image_width(seeds);
    Image s_img = create_image_from_image2d(seeds, 4);
    __global uint* seeds_ptr = (__global uint*)s_img.ptr;
    seeds_ptr = seeds_ptr + coord.x + coord.y * width;
    uint4 key = vload4(0, seeds_ptr);

    uint4 ctr = (uint4)(0);
    float4 result = 0;

    width = get_image_width(output);
    Image o_img = create_image_from_image2d(output, 4);
    coord.x = gidx * stride + width * coord.y;
    __global float* output_ptr = (__global float*)o_img.ptr;
    output_ptr += coord.x;

    for(int i = 0; i < iter; i++)
    {
        ctr = philox4x32_R_10(ctr, key);
        result = convert_float4(ctr) * re_rand_max;
        vstore4(result, i, output_ptr);
    }
}

#define logE    (1.44269502f)
float4 eltwise_unary_exp(float4 x)
{
    x *= logE;
    x = exp2(x);
    return x;
}
// N times of 8
// x dim = 1
__kernel void random_multinomial_cdf_F16
    (
    __read_only  image2d_t  input,
    __write_only image2d_t  output
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int4 coord = (int4)(gidx, gidy, 0, 0);

    vxc_half8 maxData, data;
    vxc_short8 src0;
    float4 dst0 = 0, dst1 = 0;
    float4 one = (float4)(1, 1, 1, 1);
    float tmp = 0;

    int class_max_stride = get_image_width(input);
    int offset = gidy * class_max_stride;
    Image o_img = create_image_from_image2d(output, 4);
    __global float* output_ptr = (__global float*)o_img.ptr;
    __global float* cdfPtr = output_ptr + offset;

    VXC_ReadImage(maxData, input, coord, VXC_5BITOFFSET_XY(0, 0),\
            VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    coord.x += 8;
    for(int i = 1; i < class_max_iter; i++)
    {
        VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),\
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord.x += 8;
        _viv_asm(COPY, data, src0, 16);

        VXC_VertMax3_Half(maxData, maxData, maxData, data, VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
    }
    VXC_HorzMax3_Half(maxData, maxData, VXC_MODIFIER(0, 5, 0,VXC_RM_TowardZero, 0));
    VXC_DP2x8(maxData, maxData, maxData, VXC_MODIFIER(0, 2, 0, VXC_RM_TowardZero, 0), uniPackMaxData_2x8);
    VXC_HorzMax3_Half(maxData, maxData, VXC_MODIFIER(0, 0, 0,VXC_RM_TowardZero, 0));

    coord.x = 0;
    for(int i = 0; i < class_max_iter; i++)
    {
        float4 val0, val1;
        VXC_ReadImage(src0, input, coord, VXC_5BITOFFSET_XY(0, 0),\
                VXC_MODIFIER(0, 7, 0, VXC_RM_TowardZero, 0));
        coord.x += 8;
        _viv_asm(COPY, data, src0, 16);
        VXC_DP4x4(val0, data, maxData, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGetSubData0to3_4x4);
        VXC_DP4x4(val1, data, maxData, VXC_MODIFIER(0, 3, 0, VXC_RM_TowardZero, 0), uniGetSubData4to7_4x4);
        val0 = eltwise_unary_exp(val0);
        val1 = eltwise_unary_exp(val1);
        val0.x += dst1.w;
        dst0 = (float4)(val0.x, (val0.x + val0.y), dot(val0, (float4)(1, 1, 1, 0)), dot(val0, one));
        val1.x += dst0.w;
        dst0 = (float4)(val1.x, (val1.x + val1.y), dot(val1, (float4)(1, 1, 1, 0)), dot(val1, one));
        vstore4(dst0, 0, cdfPtr);
        vstore4(dst1, 1, cdfPtr);
        cdfPtr += 8;
    }
}

__kernel void random_multinomial_cdf_F32
    (
    __read_only  image2d_t input,
    __write_only image2d_t output
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int4 coord = (int4)(gidx, gidy, 0, 0);

    vxc_float4 src0, data;
    float maxData0 = FLT_MIN, maxData1 = FLT_MIN;
    uint4 ctr = (uint4)(0);
    float4 dst = 0;
    float4 one = (float4)(1, 1, 1, 1);

    int class_max_stride = get_image_width(input);
    float tmp = 0;
    int offset = gidy * class_max_stride;
    Image o_img = create_image_from_image2d(output, 4);
    __global float* output_ptr = (__global float*)o_img.ptr;
    __global float* cdfPtr = output_ptr + offset;

    int width = get_image_width(input);
    Image i_img = create_image_from_image2d(input, 4);
    __global float* input_ptr = (__global float*)i_img.ptr;
    input_ptr = input_ptr + coord.x + coord.y * width;

    float4 maxVal = vload4(0, input_ptr);
    for(int i = 1; i < class_max_iter; i++)
    {
        src0 = vload4(i, input_ptr);

        maxVal = maxVal > src0 ? maxVal : src0;
    }
    maxVal.xy = maxVal.xy > maxVal.zw ? maxVal.xy : maxVal.zw;
    maxData0 = maxVal.x > maxVal.y ? maxVal.x : maxVal.y;

    float4 maxData = (float4)(maxData0, maxData0, maxData0, maxData0);
    for(int i = 0; i < class_max_iter; i++)
    {
        float4 val;
        src0 = vload4(i, input_ptr);
        data = src0 - maxData;
        val = eltwise_unary_exp(data);
        val.x += dst.w;
        dst = (float4)(val.x, (val.x + val.y), dot(val, (float4)(1, 1, 1, 0)), dot(val, one));
        vstore4(dst, i, cdfPtr);
    }
}

uint upper_bound(float* a, int n, float x)
{
    uint l = 0;
    uint h = n;
    while (l < h) {
        int mid = (l + h) >> 1;
        if (x >= a[mid]) {
            l = mid + 1;
        } else {
            h = mid;
        }
    }
    return l;
}

// one thread calculate 4
__kernel void random_multinomial
    (
    __read_only image2d_t   randoms,
    __read_only image2d_t   cdfs,
   __write_only image2d_t   output,
                int         class_size
    )
{
    int gidx = get_global_id(0);
    int gidy = get_global_id(1);
    int4 coord = (int4)(gidx, gidy, 0, 0);

    int class_max_stride = get_image_width(cdfs);
    int offset = gidy * class_max_stride;
    Image cdf_img = create_image_from_image2d(cdfs, 4);
    __global float* cdf_ptr = (__global float*)cdf_img.ptr;
    __global float* cdfPtr = cdf_ptr + offset;

    int width = get_image_width(randoms);
    offset = coord.x + coord.y * width;
    Image r_img = create_image_from_image2d(randoms, 4);
    __global float* randoms_ptr = (__global float*)r_img.ptr;
    randoms_ptr = randoms_ptr + offset;

    width = get_image_width(output);
    offset = coord.x + coord.y * width;
    Image o_img = create_image_from_image2d(output, 4);
    __global uint* output_ptr = (__global uint*)o_img.ptr;
    output_ptr = output_ptr + offset;

    float4 ran = vload4(0, randoms_ptr);
    float total = cdfPtr[class_size - 1];
    float4 target = ran * total;

    uint4 out_class = (uint4)(0);
    out_class.x = upper_bound(cdfPtr, class_size, target.x);
    out_class.y = upper_bound(cdfPtr, class_size, target.y);
    out_class.z = upper_bound(cdfPtr, class_size, target.z);
    out_class.w = upper_bound(cdfPtr, class_size, target.w);

    vstore4(out_class, 0, output_ptr);
}

