
/* 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. */

__kernel void fcl_compute_nuisance_jacobian_ordered_choice(
        __global fcl_fp16 *data,
        __global fcl_fp16 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        int slices,
        int nuisance_offset,
        __global fcl_fp *jacobian_out,
        __global fcl_fp16 *nuisance_jacobian_prefix,
        __global fcl_fp16 *nuisance_jacobian1_prefix) {
    int j = get_global_id(0);
    
    int i;
    fcl_fp result = 0.f;
    for (i=0; i<slices; i++) {
        int16 actual = int_array_elements16(dep_vals_reverse_lookup,
                convert_fp16_to_int16(data[dependent_col*slices+i]));
        result += dot16(fweights[i],
                select16(select16(0.0, nuisance_jacobian_prefix[i], actual-1 == j),
                    nuisance_jacobian1_prefix[i], actual-2 == j));
    }
    jacobian_out[nuisance_offset+j] = result;
}

__kernel void fcl_compute_diagonal_nuisance_hessian_ordered_choice(
        __global fcl_fp16 *data,
        __global fcl_fp16 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        int slices,
        int hessian_cols_count,
        int nuisance_offset,
        __global fcl_fp *hessian_out,
        __global fcl_fp16 *nuisance_hessian_prefix,
        __global fcl_fp16 *nuisance_hessian1_prefix,
        __global fcl_fp16 *nuisance_hessian_cross_prefix) 
{
    int j = get_global_id(0);
    int i;
    fcl_fp result = 0.f;
    for (i=0; i<slices; i++) {
        int16 actual = int_array_elements16(dep_vals_reverse_lookup,
                convert_fp16_to_int16(data[dependent_col*slices+i]));
        result += dot16(fweights[i],
                select16(select16(0.0, nuisance_hessian_prefix[i], actual-1 == j),
                    nuisance_hessian1_prefix[i], actual-2 == j));
    }
    hessian_out[hessian_cols_count*(nuisance_offset+j)+nuisance_offset+j] = result;
}

__kernel void fcl_compute_offdiagonal_nuisance_hessian_ordered_choice(
        __global fcl_fp16 *data,
        __global fcl_fp16 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        int slices,
        int hessian_cols_count,
        int nuisance_offset,
        __global fcl_fp *hessian_out,
        __global fcl_fp16 *nuisance_hessian_prefix,
        __global fcl_fp16 *nuisance_hessian1_prefix,
        __global fcl_fp16 *nuisance_hessian_cross_prefix) {
    
    int j = get_global_id(0);
    int i;
    fcl_fp result = 0.f;
    for (i=0; i<slices; i++) {
        int16 actual = int_array_elements16(dep_vals_reverse_lookup,
                convert_fp16_to_int16(data[dependent_col*slices+i]));
        result += dot16(fweights[i],
                select16(0.0, nuisance_hessian_cross_prefix[i], actual-2 == j));
    }
    hessian_out[hessian_cols_count*(nuisance_offset+j)+(nuisance_offset+j+1)] = result;
    hessian_out[hessian_cols_count*(nuisance_offset+j+1)+(nuisance_offset+j)] = result;
}

__kernel void fcl_compute_cross_hessian_ordered_choice(
        __global fcl_fp16 *data,
        __global fcl_fp16 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        __global int *independent_cols,
        int slices,
        __global int *id2col,
        int hessian_cols_count,
        int nuisance_offset,
        __global fcl_fp *hessian_out,
        __global fcl_fp16 *cross_hessian_prefix,
        __global fcl_fp16 *cross_hessian1_prefix) 
{
    int k = get_global_id(0);
    int j = get_global_id(1);
    int idx = id2col[k];
    int i;
    __global fcl_fp16 *col = &data[independent_cols[idx]*slices];

    fcl_fp result = 0.f;
    for (i=0; i<slices; i++) {
        int16 actual = int_array_elements16(dep_vals_reverse_lookup,
                convert_fp16_to_int16(data[dependent_col*slices+i]));
        result += dot16(fweights[i] * col[i],
                select16(select16(0.0, cross_hessian_prefix[i], actual-1 == j),
                    cross_hessian1_prefix[i], actual-2 == j));
    }
    hessian_out[hessian_cols_count*k+nuisance_offset+j] = result;
    hessian_out[hessian_cols_count*(nuisance_offset+j)+k] = result;
}

__kernel void fcl_compute_cross_opg_robust_ordered_choice(
        __global fcl_fp16 *data,
        __global fcl_fp16 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        __global int *independent_cols,
        int slices,
        __global int *id2col,
        __global fcl_fp16 *core_jacobian_prefix,
        __global fcl_fp16 *nuisance_jacobian_prefix,
        __global fcl_fp16 *nuisance_jacobian1_prefix,
        __global fcl_fp *opg_out,
        int hessian_cols_count,
        int nuisance_offset) 
{
    int k = get_global_id(0);
    int j = get_global_id(1);
    int idx = id2col[k];
    int i;

	__global fcl_fp16 *col1 = &data[independent_cols[idx]*slices];

    fcl_fp result = 0.f;
    for (i=0; i<slices; i++) {
        int16 actual = int_array_elements16(dep_vals_reverse_lookup,
                convert_fp16_to_int16(data[dependent_col*slices+i]));
        fcl_fp16 residual = core_jacobian_prefix[i];
        result += dot16(fweights[i] * col1[i] * residual,
                select16(select16(0.0, nuisance_jacobian1_prefix[i], actual-2 == j), 
                    nuisance_jacobian_prefix[i], actual-1 == j));
    }

    opg_out[hessian_cols_count*k+nuisance_offset+j] = result;
    opg_out[hessian_cols_count*(nuisance_offset+j)+k] = result;
}

__kernel void fcl_compute_diagonal_nuisance_opg_robust_ordered_choice(
        __global fcl_fp16 *data,
        __global fcl_fp16 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        __global int *independent_cols,
        int slices,
        __global int *id2col,
        __global fcl_fp16 *core_jacobian_prefix,
        __global fcl_fp16 *nuisance_jacobian_prefix,
        __global fcl_fp16 *nuisance_jacobian1_prefix,
        __global fcl_fp *opg_out,
        int hessian_cols_count,
        int nuisance_offset) {
    int j = get_global_id(0);

    fcl_fp result = 0.f;
    int i;
    for (i=0; i<slices; i++) {
        int16 actual = int_array_elements16(dep_vals_reverse_lookup,
                convert_fp16_to_int16(data[dependent_col*slices+i]));
        fcl_fp16 value = select16(select16(0.0, nuisance_jacobian_prefix[i], actual-1 == j),
                nuisance_jacobian1_prefix[i], actual-2 == j);
        result += dot16(fweights[i] * value, value);
    }
    opg_out[hessian_cols_count*(nuisance_offset+j)+nuisance_offset+j] = result;
}

__kernel void fcl_compute_offdiagonal_nuisance_opg_robust_ordered_choice(
        __global fcl_fp16 *data,
        __global fcl_fp16 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        __global int *independent_cols,
        int slices,
        __global int *id2col,
        __global fcl_fp16 *core_jacobian_prefix,
        __global fcl_fp16 *nuisance_jacobian_prefix,
        __global fcl_fp16 *nuisance_jacobian1_prefix,
        __global fcl_fp *opg_out,
        int hessian_cols_count,
        int nuisance_offset
        ) {
    int j = get_global_id(0);

    fcl_fp result = 0.f;
    int i;
    for (i=0; i<slices; i++) {
        int16 actual = int_array_elements16(dep_vals_reverse_lookup,
                convert_fp16_to_int16(data[dependent_col*slices+i]));
        result += dot16(fweights[i],
                select16(0.0, nuisance_jacobian_prefix[i] * nuisance_jacobian1_prefix[i], actual-2 == j));
    }
    opg_out[hessian_cols_count*(nuisance_offset+j)+(nuisance_offset+j+1)] = result;
    opg_out[hessian_cols_count*(nuisance_offset+j+1)+(nuisance_offset+j)] = result;
}

__kernel void fcl_compute_nuisance_opg_cluster_pass1_ordered_choice(
        __global fcl_fp4 *data,
        __global fcl_fp4 *fweights,
        __global int *independent_cols,
        int slivers,
        __global int *id2col,
        __global int *clusters,
        __global int *panel_sizes,
        __global int *panel_offsets,
        __global fcl_fp4 *jacobian_prefix,
        __global fcl_fp *cluster_grad,
        __global fcl_fp4 *nuisance_jacobian_prefix,
        __global fcl_fp4 *nuisance_jacobian1_prefix,
        int nuisance_offset,
        int dependent_col,
        __global int *dep_vals_reverse_lookup) {
	int j = get_global_id(0);
    int i = get_global_id(1);
    int cluster_count = get_global_size(1);
    int panel_idx = clusters[i];
    
    fcl_fp result = 0.f;
    int t;
    int last = panel_offsets[panel_idx] + panel_sizes[panel_idx];
	for (t=panel_offsets[panel_idx]; t<last; t++) {
        int4 actual = int_array_elements4(dep_vals_reverse_lookup,
                convert_fp4_to_int4(data[dependent_col*slivers+t]));
		result += dot(fweights[t], select4(select4(0.f, nuisance_jacobian1_prefix[t], actual-2 == j),
                nuisance_jacobian_prefix[t], actual-1 == j));
	}
    cluster_grad[(nuisance_offset+j)*cluster_count+i] = result;
}

__kernel void fcl_compute_fixed_effects_nuisance_cross_hessian_ordered_choice(
        __global fcl_fp4 *data,
        __global fcl_fp4 *fweights,
        int dependent_col,
        __global int *dep_vals_reverse_lookup,
        int slivers,
        __global int *clusters,
        __global int *panel_sizes,
        __global int *panel_offsets,
        __global fcl_fp *hessian_prefix_sum,
        __global fcl_fp4 *cross_hessian_prefix,
        __global fcl_fp4 *cross_hessian1_prefix,
        int nuisance_offset,
        __global fcl_fp *out_fe_work,
        __global fcl_fp *out_fe_mean) 
{
    int j = get_global_id(0);
    int i = get_global_id(1);
    int panel_idx = clusters[i];

    int cluster_count = get_global_size(1);

    int t;
    int last = panel_offsets[panel_idx] + panel_sizes[panel_idx];
    fcl_fp result = 0.f;
    for (t=panel_offsets[panel_idx]; t<last; t++) {
        int4 actual = int_array_elements4(dep_vals_reverse_lookup,
                convert_fp4_to_int4(data[dependent_col*slivers+t]));
        result += dot(fweights[t],
                select4(select4(0.0, cross_hessian_prefix[t], actual-1 == j),
                    cross_hessian1_prefix[t], actual-2 == j));
    }

    out_fe_work[cluster_count*(nuisance_offset+j) + i] = result;
    out_fe_mean[cluster_count*(nuisance_offset+j) + i] = result / hessian_prefix_sum[i];
}
