
/* Copyright 2012 Evan Miller */
/* All rights reserved. */

#ifdef cl_khr_fp64
typedef double   fcl_fp;
typedef double4  fcl_fp4;
typedef double16 fcl_fp16;
#else
typedef float    fcl_fp;
typedef float4   fcl_fp4;
typedef float16  fcl_fp16;
#endif

#ifndef M_SQRT1_2
#define M_SQRT1_2 M_SQRT1_2_F
#endif

#ifndef M_LN2
#define M_LN2 M_LN2_F
#endif

#ifdef cl_khr_fp64
#define FCL_1_OVER_SQRT_2_PI 0.3989422804014326
#else
#define FCL_1_OVER_SQRT_2_PI 0.3989422804014326f
#endif

#ifdef cl_khr_fp64
#define FCL_DIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX   8.5
#define FCL_DIGAMMA_MAX_ARG_FOR_LINEAR_APPROX       0.00001
#define FCL_DIGAMMA_D1                             -0.5772156649
#define FCL_DIGAMMA_S3                              0.08333333333
#define FCL_DIGAMMA_S4                              0.0083333333333
#define FCL_DIGAMMA_S5                              0.003968253968
#else
#define FCL_DIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX   8.5f
#define FCL_DIGAMMA_MAX_ARG_FOR_LINEAR_APPROX       0.00001f
#define FCL_DIGAMMA_D1                             -0.5772156649f
#define FCL_DIGAMMA_S3                              0.08333333333f
#define FCL_DIGAMMA_S4                              0.0083333333333f
#define FCL_DIGAMMA_S5                              0.003968253968f
#endif

#ifdef cl_khr_fp64
#define FCL_TRIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX  5.
#define FCL_TRIGAMMA_MAX_ARG_FOR_LINEAR_APPROX      0.0001
#define FCL_TRIGAMMA_B2                             0.1666666667
#define FCL_TRIGAMMA_B4                            -0.03333333333
#define FCL_TRIGAMMA_B6                             0.02380952381
#define FCL_TRIGAMMA_B8                            -0.03333333333
#else
#define FCL_TRIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX  5.f
#define FCL_TRIGAMMA_MAX_ARG_FOR_LINEAR_APPROX      0.0001f
#define FCL_TRIGAMMA_B2                             0.1666666667f
#define FCL_TRIGAMMA_B4                            -0.03333333333f
#define FCL_TRIGAMMA_B6                             0.02380952381f
#define FCL_TRIGAMMA_B8                            -0.03333333333f
#endif

#define fcl_slice_to_sliver(i) (4*i)

fcl_fp sum4(fcl_fp4 vec);
fcl_fp sum16(fcl_fp16 vec);
fcl_fp dot16(fcl_fp16 vec1, fcl_fp16 vec2);

fcl_fp4 select4(fcl_fp4 c, fcl_fp4 b, int4 a);
fcl_fp16 select16(fcl_fp16 c, fcl_fp16 b, int16 a);

fcl_fp fcl_logistic_cdf(fcl_fp alpha);
fcl_fp fcl_normal_cdf(fcl_fp alpha);

fcl_fp fcl_normal_pdf(fcl_fp alpha);
fcl_fp fcl_inverse_mills_ratio(fcl_fp alpha);

fcl_fp digamma(fcl_fp alpha);
fcl_fp16 digamma16(fcl_fp16 alpha);

fcl_fp trigamma(fcl_fp alpha);
fcl_fp16 trigamma16(fcl_fp16 alpha);

int4 convert_fp4_to_int4(fcl_fp4 res_fp);
int16 convert_fp16_to_int16(fcl_fp16 res_fp);

fcl_fp4 convert_int4_to_fp4(int4 res_i);
fcl_fp16 convert_int16_to_fp16(int16 res_i);

int4 int_array_elements4(__global int *array, int4 idx);
int16 int_array_elements16(__global int *array, int16 idx);
fcl_fp4 array_elements4(__global fcl_fp *array, int4 idx);
fcl_fp16 array_elements16(__global fcl_fp *array, int16 idx);

fcl_fp dot16(fcl_fp16 vec1, fcl_fp16 vec2) {
    fcl_fp result = 0.f;
    result += dot(vec1.lo.lo, vec2.lo.lo);
    result += dot(vec1.lo.hi, vec2.lo.hi);
    result += dot(vec1.hi.lo, vec2.hi.lo);
    result += dot(vec1.hi.hi, vec2.hi.hi);
    return result;
}

fcl_fp4 select4(fcl_fp4 c, fcl_fp4 b, int4 a) {
    fcl_fp4 res_fp;
    res_fp.s0 = a.s0 ? b.s0 : c.s0;
    res_fp.s1 = a.s1 ? b.s1 : c.s1;
    res_fp.s2 = a.s2 ? b.s2 : c.s2;
    res_fp.s3 = a.s3 ? b.s3 : c.s3;
    return res_fp;
}

fcl_fp16 select16(fcl_fp16 c, fcl_fp16 b, int16 a) {
    fcl_fp16 res_fp;
    res_fp.lo.lo = select4(c.lo.lo, b.lo.lo, a.lo.lo);
    res_fp.lo.hi = select4(c.lo.hi, b.lo.hi, a.lo.hi);
    res_fp.hi.lo = select4(c.hi.lo, b.hi.lo, a.hi.lo);
    res_fp.hi.hi = select4(c.hi.hi, b.hi.hi, a.hi.hi);
    return res_fp;
}

fcl_fp sum4(fcl_fp4 vec) {
    return vec.x + vec.y + vec.z + vec.w;
}

fcl_fp sum16(fcl_fp16 vec) {
    return sum4(vec.lo.lo) + sum4(vec.lo.hi) + sum4(vec.hi.lo) + sum4(vec.hi.hi);
}

fcl_fp fcl_logistic_cdf(fcl_fp alpha) {
    return 1.f / (1.f + exp(-alpha));
}

fcl_fp fcl_normal_cdf(fcl_fp alpha) {
    return 0.5f * erfc(-alpha * M_SQRT1_2);
}

fcl_fp fcl_normal_pdf(fcl_fp alpha) {
    return FCL_1_OVER_SQRT_2_PI * exp(-0.5f * alpha * alpha);
}

fcl_fp fcl_inverse_mills_ratio(fcl_fp alpha) {
    return fcl_normal_pdf(alpha) / fcl_normal_cdf(alpha);
}

fcl_fp digamma(fcl_fp alpha) {
    fcl_fp r;
    fcl_fp value = 0.f;
    fcl_fp y = alpha;

    if (alpha <= FCL_DIGAMMA_MAX_ARG_FOR_LINEAR_APPROX) {
        return FCL_DIGAMMA_D1 - 1.f / y;
    }
    while (y < FCL_DIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX) {
        value -= 1.f / y;
        y += 1.f;
    }
    value += log(y) - 0.5f / y;
    r = 1.f / (y * y);

    return value - r*(FCL_DIGAMMA_S3 - r*(FCL_DIGAMMA_S4 - r*FCL_DIGAMMA_S5));
}

fcl_fp16 digamma16(fcl_fp16 alpha) {
    fcl_fp16 r;
    fcl_fp16 y = alpha;
    fcl_fp16 value = (fcl_fp16)(0.f);
    while (any(y < FCL_DIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX)) {
        value -= 1.f / y;
        y += 1.f;
    }
    value += log(y) - 0.5f / y;
    r = 1.f / (y * y);

    return select(value - r*(FCL_DIGAMMA_S3 - r*(FCL_DIGAMMA_S4 - r*FCL_DIGAMMA_S5)), 
            FCL_DIGAMMA_D1 - 1.f / alpha, alpha <= FCL_DIGAMMA_MAX_ARG_FOR_LINEAR_APPROX);
}

fcl_fp trigamma(fcl_fp alpha) {
    fcl_fp value = 0.f;
    fcl_fp y;
    fcl_fp z = alpha;

    if (alpha <= FCL_TRIGAMMA_MAX_ARG_FOR_LINEAR_APPROX) {
        return 1.f / (alpha * alpha);
    }

    while (z < FCL_TRIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX) {
        value += 1.f / (z * z);
        z += 1.f;
    }

    y = 1.f / (z * z);

    return value + 0.5f * y + ( 1.f
            + y * ( FCL_TRIGAMMA_B2
                + y * ( FCL_TRIGAMMA_B4
                    + y * ( FCL_TRIGAMMA_B6
                        + y * FCL_TRIGAMMA_B8 )))) / z;
}

fcl_fp16 trigamma16(fcl_fp16 alpha) {
    fcl_fp16 value = 0.f;
    fcl_fp16 y;
    fcl_fp16 z = alpha;

    while (any(z < FCL_TRIGAMMA_MIN_ARG_FOR_POLYNOMIAL_APPROX)) {
        value += 1.f / (z * z);
        z += 1.f;
    }

    y = 1.f / (z * z);

    return select(value + 0.5f * y + ( 1.f
            + y * ( FCL_TRIGAMMA_B2
                + y * ( FCL_TRIGAMMA_B4
                    + y * ( FCL_TRIGAMMA_B6
                        + y * FCL_TRIGAMMA_B8 )))) / z,
                        1.f / (alpha * alpha), alpha <= FCL_TRIGAMMA_MAX_ARG_FOR_LINEAR_APPROX);
}

int4 convert_fp4_to_int4(fcl_fp4 res_fp) {
    int4 res_i;
    res_i.s0 = res_fp.s0;
    res_i.s1 = res_fp.s1;
    res_i.s2 = res_fp.s2;
    res_i.s3 = res_fp.s3;
    return res_i;
}

int16 convert_fp16_to_int16(fcl_fp16 res_fp) {
    int16 res_i;
    res_i.lo.lo = convert_fp4_to_int4(res_fp.lo.lo);
    res_i.lo.hi = convert_fp4_to_int4(res_fp.lo.hi);
    res_i.hi.lo = convert_fp4_to_int4(res_fp.hi.lo);
    res_i.hi.hi = convert_fp4_to_int4(res_fp.hi.hi);
    return res_i;
}

fcl_fp4 convert_int4_to_fp4(int4 res_i) {
    fcl_fp4 res_fp;
    res_fp.s0 = res_i.s0;
    res_fp.s1 = res_i.s1;
    res_fp.s2 = res_i.s2;
    res_fp.s3 = res_i.s3;
    return res_fp;
}

fcl_fp16 convert_int16_to_fp16(int16 res_i) {
    fcl_fp16 res_fp;
    res_fp.lo.lo = convert_int4_to_fp4(res_i.lo.lo);
    res_fp.lo.hi = convert_int4_to_fp4(res_i.lo.hi);
    res_fp.hi.lo = convert_int4_to_fp4(res_i.hi.lo);
    res_fp.hi.hi = convert_int4_to_fp4(res_i.hi.hi);
    return res_fp;
}

int4 int_array_elements4(__global int *array, int4 idx) {
    int4 res_i;
    res_i.s0 = array[idx.s0];
    res_i.s1 = array[idx.s1];
    res_i.s2 = array[idx.s2];
    res_i.s3 = array[idx.s3];
    return res_i;
}

int16 int_array_elements16(__global int *array, int16 idx) {
    int16 res_i;
    res_i.lo.lo = int_array_elements4(array, idx.lo.lo);
    res_i.lo.hi = int_array_elements4(array, idx.lo.hi);
    res_i.hi.lo = int_array_elements4(array, idx.hi.lo);
    res_i.hi.hi = int_array_elements4(array, idx.hi.hi);
    return res_i;
}

/* Copyright 2012 Evan Miller */
/* All rights reserved. */

/* Linear and weighted linear */

__kernel void fcl_prefix_linear(
        __global fcl_fp *data,
        int dependent_col,
        __global fcl_fp *fitted,
        __global fcl_fp *jacobian_prefix_out,
        __global fcl_fp *hessian_prefix_out
        ) {
	int i = get_global_id(0);
    int slices = get_global_size(0);

    fcl_fp actual = data[dependent_col*slices+i];
    fcl_fp xb = fitted[i];

	jacobian_prefix_out[i] = actual - xb;
	hessian_prefix_out[i] = -1.f;
}

__kernel void fcl_prefix_log_linear(
        __global fcl_fp *data,
        int dependent_col,
        __global fcl_fp *fitted,
        __global fcl_fp *jacobian_prefix_out,
        __global fcl_fp *hessian_prefix_out
        ) {
	int i = get_global_id(0);
    int slices = get_global_size(0);

    fcl_fp actual = data[dependent_col*slices+i];
    fcl_fp xb = fitted[i];

	jacobian_prefix_out[i] = actual > 0.f ? log(actual) - xb : 0.f;
	hessian_prefix_out[i] = -1.f;
}

__kernel void fcl_prefix_weighted_linear(
        __global fcl_fp *data,
        int dependent_col,
        __global fcl_fp *fitted,
        __global fcl_fp *jacobian_prefix_out,
        __global fcl_fp *hessian_prefix_out,
        __global fcl_fp *weights
        ) {
	int i = get_global_id(0);
    int slices = get_global_size(0);

    fcl_fp actual = data[dependent_col*slices+i];
    fcl_fp weight = weights[i];
    fcl_fp xb = fitted[i];

	jacobian_prefix_out[i] = (actual - xb) * weight;
	hessian_prefix_out[i] = - weight;
}

__kernel void fcl_prefix_weighted_log_linear(
        __global fcl_fp *data,
        int dependent_col,
        __global fcl_fp *fitted,
        __global fcl_fp *jacobian_prefix_out,
        __global fcl_fp *hessian_prefix_out,
        __global fcl_fp *weights) {
	int i = get_global_id(0);
    int slices = get_global_size(0);

    fcl_fp actual = data[dependent_col*slices+i];
    fcl_fp log_actual = actual > 0.f ? log(actual) : 0.f;
    fcl_fp weight = weights[i];
    fcl_fp xb = fitted[i];

	jacobian_prefix_out[i] = (log_actual - xb) * weight;
	hessian_prefix_out[i] = - weight;
}

__kernel void fcl_prefix_weighted_linear_fe(
                                            __global fcl_fp16 *data,
                                            int dependent_col,
                                            __global fcl_fp16 *fitted,
                                            __global fcl_fp16 *jacobian_prefix_out,
                                            __global fcl_fp16 *hessian_prefix_out,
                                            __global fcl_fp *fe_estimates,
                                            __global int *sliver2panel,
                                            __global fcl_fp *weights
                                         ) {
    int i = get_global_id(0);
    int slices = get_global_size(0);
    
    int sliver = fcl_slice_to_sliver(i);
    
    fcl_fp16 actual = data[dependent_col*slices+i];
    fcl_fp16 weight = (fcl_fp16)(
            (fcl_fp4)(weights[sliver2panel[sliver]]),
            (fcl_fp4)(weights[sliver2panel[sliver+1]]), 
            (fcl_fp4)(weights[sliver2panel[sliver+2]]), 
            (fcl_fp4)(weights[sliver2panel[sliver+3]]));
    
	jacobian_prefix_out[i] = (actual - fitted[i]) * weight;
	hessian_prefix_out[i] = - weight * weight;
}

__kernel void fcl_prefix_weighted_log_linear_fe(
        __global fcl_fp16 *data,
        int dependent_col,
        __global fcl_fp16 *fitted,
        __global fcl_fp16 *jacobian_prefix_out,
        __global fcl_fp16 *hessian_prefix_out,
        __global fcl_fp *fe_estimates,
        __global int *sliver2panel,
        __global fcl_fp *weights) {
    int i = get_global_id(0);
    int slices = get_global_size(0);
    
    int sliver = 4*i;
    
    fcl_fp16 actual = data[dependent_col*slices+i];
    fcl_fp16 log_actual = select(0.f, log(actual), actual > 0.f);
    fcl_fp16 weight = (fcl_fp16)(
            (fcl_fp4)(weights[sliver2panel[sliver]]),
            (fcl_fp4)(weights[sliver2panel[sliver+1]]), 
            (fcl_fp4)(weights[sliver2panel[sliver+2]]), 
            (fcl_fp4)(weights[sliver2panel[sliver+3]]));
    
	jacobian_prefix_out[i] = (log_actual - fitted[i]) * weight;
	hessian_prefix_out[i] = - weight * weight;
}

__kernel void fcl_predict_log_linear(
        __global fcl_fp *fitted,
        __global fcl_fp *predicted_out) {
	int i = get_global_id(0);
    fcl_fp xb = fitted[i];
    predicted_out[i] = exp(xb);
}
