
/* 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_jacobian_multinomial_logit(
		__global fcl_fp4 *data,
		__global fcl_fp4 *fweights,
        __global int *independent_cols,
		int slices,
        __global int *id2col,
		__global fcl_fp4 *jacobian_prefix,
		__global fcl_fp *jacobian_out,
        __global int *id2dep_val,
        int dep_vals_count
		) {
    int d = get_global_id(0);
    
	int j = id2dep_val[d];
	int k = id2col[d];
    
    int slivers = fcl_slice_to_sliver(slices);
        
	int i;

	__global fcl_fp4 *col = &data[independent_cols[k]*slivers];

	fcl_fp result = 0.f;
    int prefix_offset = j;
	for (i=0; i<slivers; i++) {
        result += dot(jacobian_prefix[prefix_offset], fweights[i] * col[i]);
        prefix_offset += dep_vals_count;
	}
	jacobian_out[d] = result;
}

__kernel void fcl_compute_hessian_multinomial_logit(
		__global fcl_fp4 *data,
		__global fcl_fp4 *fweights,
        __global int *independent_cols,
		int slices,
        __global int *id2col,
		__global fcl_fp4 *hessian_prefix,
		__global fcl_fp *hessian_out,
        int row_size,
		__global int *id2dep_val,
		int dep_vals_count
    ) {
	int d1 = get_global_id(0);
	int d2 = get_global_id(1);

	int j1 = id2dep_val[d1];
	int k1 = id2col[d1];
	int j2 = id2dep_val[d2];
	int k2 = id2col[d2];
    
    int slivers = fcl_slice_to_sliver(slices);
            
	__global fcl_fp4 *col1 = &data[independent_cols[k1]*slivers];
	__global fcl_fp4 *col2 = &data[independent_cols[k2]*slivers];

	int i;
	fcl_fp result = 0.f;
	int prefix_offset = j1*dep_vals_count+j2;

    if (d2 > d1)
        return;
    
	for (i=0; i<slivers; i++) {
        result += dot(fweights[i] * hessian_prefix[prefix_offset], col1[i] * col2[i]);

		prefix_offset += dep_vals_count * dep_vals_count;
	}

	hessian_out[d1*row_size+d2] = result;
	hessian_out[d2*row_size+d1] = result;
}

__kernel void fcl_compute_opg_robust_multinomial_logit(
		__global fcl_fp4 *data,
		__global fcl_fp4 *fweights,
        __global int *independent_cols,
		int slices,
        __global int *id2col,
        int row_size,
		__global fcl_fp4 *jacobian_prefix,
		__global fcl_fp *opg_out,
        __global int *id2dep_val,        
		int dep_vals_count
		) {
	int d1 = get_global_id(0);
	int d2 = get_global_id(1);

    int slivers = fcl_slice_to_sliver(slices);

	int j1 = id2dep_val[d1];
	int k1 = id2col[d1];
	int j2 = id2dep_val[d2];
	int k2 = id2col[d2];
    
	__global fcl_fp4 *col1 = &data[independent_cols[k1]*slivers];
	__global fcl_fp4 *col2 = &data[independent_cols[k2]*slivers];

	int i;
	fcl_fp result = 0.f;
    int prefix1_offset = j1;
    int prefix2_offset = j2;

    if (d2 > d1)
        return;
    
	for (i=0; i<slivers; i++) {
        result += dot(fweights[i] * jacobian_prefix[prefix1_offset] * col1[i], 
                jacobian_prefix[prefix2_offset] * col2[i]);
        prefix1_offset += dep_vals_count;
        prefix2_offset += dep_vals_count;
	}
	opg_out[d1*row_size+d2] = result;
	opg_out[d2*row_size+d1] = result;
}

__kernel void fcl_compute_opg_cluster_pass1_multinomial_logit(
        __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 int *id2dep_val,
        int dep_vals_count
		) {
    int d = get_global_id(0);
    int j = id2dep_val[d];
	int k = id2col[d];
    int i = get_global_id(1);
    int cluster_count = get_global_size(1);
    int panel_idx = clusters[i];
	__global fcl_fp4 *col = &data[independent_cols[k]*slivers];
	int offset = d*cluster_count+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++) {
		result += dot(fweights[t] * col[t], jacobian_prefix[t*dep_vals_count+j]);
	}
    cluster_grad[offset] = result;
}

__kernel void fcl_compute_fixed_effects_jacobian_prefix_sum_multinomial_logit(
        __global fcl_fp4 *jacobian_prefix,
        __global fcl_fp4 *fweights,
        __global int *clusters,
        __global int *panel_sizes,
        __global int *panel_offsets,
        __global fcl_fp *fe_sum_out,
        int dep_vals_count,
        __global int *free_dep_val2dep_val) {
    int i=get_global_id(0);
    int panel_idx = clusters[i];
    int f_j=get_global_id(1);
    int free_dep_vals_count = get_global_size(1);
    int j = free_dep_val2dep_val[f_j];
    
    fcl_fp result = 0.f;
    int last = panel_offsets[panel_idx] + panel_sizes[panel_idx];
    int t;
    for (t=panel_offsets[panel_idx]; t<last; t++) {
        result += dot(fweights[t], jacobian_prefix[t*dep_vals_count+j]);
    }
    
    fe_sum_out[i*free_dep_vals_count+f_j] = result;
}

__kernel void fcl_compute_fixed_effects_diagonal_hessian_multinomial_logit(
        __global fcl_fp4 *hessian_prefix,
        __global fcl_fp4 *fweights,
        __global int *clusters,
        __global int *panel_sizes,
        __global int *panel_offsets,
        __global fcl_fp *hessian_diagonal_out,
        int dep_vals_count,
        __global int *free_dep_val2dep_val)
{
    int f_j1 = get_global_id(0);
    int f_j2 = get_global_id(1);
    int free_dep_vals_count = get_global_size(0);
    int j1 = free_dep_val2dep_val[f_j1];
    int j2 = free_dep_val2dep_val[f_j2];
    int i = get_global_id(2);
    int panel_idx = clusters[i];
    int offset = panel_offsets[panel_idx]*dep_vals_count*dep_vals_count + j1*dep_vals_count + j2;
    int t;
    fcl_fp result = 0.f;
    int last = panel_offsets[panel_idx] + panel_sizes[panel_idx];
    for (t=panel_offsets[panel_idx]; t<last; t++) {
        result += dot(fweights[t], hessian_prefix[offset]);
        offset += dep_vals_count * dep_vals_count;
    }

    hessian_diagonal_out[i*free_dep_vals_count*free_dep_vals_count+f_j1*free_dep_vals_count+f_j2] = result;
}

__kernel void fcl_compute_fixed_effects_cross_hessian_multinomial_logit(
        __global fcl_fp4 *data,
        __global fcl_fp4 *fweights,
        __global int *independent_cols,
        int independent_cols_count,
        int slivers,
        __global int *id2col,
        __global int *clusters,
        __global int *panel_sizes,
        __global int *panel_offsets,
        __global fcl_fp4 *hessian_prefix,
        __global fcl_fp *hessian_cross,
        __global int *id2dep_val,
        int dep_vals_count,
        __global int *free_dep_val2dep_val)
{
    int d = get_global_id(0);
    int j1 = id2dep_val[d];
	int k1 = id2col[d];
    
    int f_j2 = get_global_id(1);
    int free_dep_vals_count = get_global_size(1);
    int j2 = free_dep_val2dep_val[f_j2];
    
    int i = get_global_id(2);
    int panel_idx = clusters[i];
    int fe_count = get_global_size(2) * free_dep_vals_count;
    
    __global fcl_fp4 *col1 = &data[independent_cols[k1]*slivers];
    
    int offset = panel_offsets[panel_idx]*dep_vals_count*dep_vals_count+j1*dep_vals_count+j2;
    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++) {
        result += dot(fweights[t] * col1[t], hessian_prefix[offset]);
        offset += dep_vals_count * dep_vals_count;
    }
    hessian_cross[d*fe_count+i*free_dep_vals_count+f_j2] = result;
}

__kernel void fcl_compute_fixed_effects_multiply_cross_by_diagonal_multinomial_logit(
        int free_dep_vals_count,
        __global fcl_fp *hessian_cross,
        __global fcl_fp *hessian_diagonal,
        __global fcl_fp *hessian_cross_out)
{
    int d = get_global_id(0);
    int i = get_global_id(1);
    int fe_count = get_global_size(1) * free_dep_vals_count;
    
    __global fcl_fp *hessian_row = hessian_cross + d*fe_count + i * free_dep_vals_count;
    __global fcl_fp *hessian_box = hessian_diagonal + i * free_dep_vals_count * free_dep_vals_count;
    __global fcl_fp *hessian_new = hessian_cross_out + d*fe_count + i * free_dep_vals_count;
    
    int j2, j3;
    for (j2=0; j2<free_dep_vals_count; j2++) {
        fcl_fp result = 0.f;
        for (j3=0; j3<free_dep_vals_count; j3++) {
            result += hessian_row[j3] * hessian_box[j2*free_dep_vals_count+j3];
        }
        hessian_new[j2] = result;
    }
}

__kernel void fcl_compute_fixed_effects_delta_hessian_multinomial_logit(
        int fe_count,
        __global fcl_fp *hessian_cross,
        __global fcl_fp *hessian_cross2,
        __global fcl_fp *hessian)
{
    
    int d1 = get_global_id(0);
    int d2 = get_global_id(1);
    
    int row_size = get_global_size(0);
    
    __global fcl_fp *hessian_col1 = hessian_cross + fe_count*d1;
    __global fcl_fp *hessian_col2 = hessian_cross2 + fe_count*d2;
    
    int i;
    fcl_fp delta = 0.f;
    for (i=0; i<fe_count; i++) {
        delta += hessian_col1[i] * hessian_col2[i];
    }
    
    hessian[d1*row_size+d2] -= delta;
}

__kernel void fcl_compute_fixed_effects_delta_jacobian_multinomial_logit(
                                                                    int fe_count,
                                                                    __global fcl_fp *hessian_cross2,
                                                                    __global fcl_fp *jacobian_fe_sum,
                                                                    __global fcl_fp *jacobian
) {
    int d1 = get_global_id(0);
    
    __global fcl_fp *hessian_col = hessian_cross2 + fe_count*d1;
    
    int i;
    fcl_fp delta = 0.f;
    for (i=0; i<fe_count; i++) {
        delta += hessian_col[i] * jacobian_fe_sum[i];
    }
    
    jacobian[d1] -= delta;
}

__kernel void fcl_compute_fixed_effects_delta_multinomial_logit(
        int free_dep_vals_count,
        int hessian_cols_count,
        int fe_count,
        __global fcl_fp *jacobian_fe,
        __global fcl_fp *hessian_diagonal,
        __global fcl_fp *hessian_cross2,
        __global fcl_fp *delta_beta,
        __global fcl_fp *delta_fe
) {
    int i = get_global_id(0);
    int cluster_count = get_global_size(0);
    
    __global fcl_fp *hessian_box = hessian_diagonal + i * free_dep_vals_count * free_dep_vals_count;
    
    int cols = hessian_cols_count;
    int f_j, f_j2, d;
    for (f_j=0; f_j<free_dep_vals_count; f_j++) {
        fcl_fp result1 = 0.f, result2 = 0.f;
        for (f_j2=0; f_j2<free_dep_vals_count; f_j2++) {
            result1 += hessian_box[f_j*free_dep_vals_count+f_j2] * jacobian_fe[i*free_dep_vals_count+f_j2];
        }
        for (d=0; d<cols; d++) {
            result2 += hessian_cross2[fe_count*d+i*free_dep_vals_count+f_j] * delta_beta[d];
        }
        delta_fe[f_j*cluster_count+i] = result1 - result2;
    }
}
