functions {
// Variance masking function
  vector STAN_var_mask(vector x, real a){
    return( inv(1+exp(-a*x)) );
  }

  // Expand a vector
  vector STAN_expand(vector v, data array[] int idx_expand){
    int L = num_elements(v);
    vector[L+1] v_add0 = rep_vector(0.0, L+1);
    v_add0[2:(L+1)] = v;
    return(v_add0[idx_expand]);
  }

    // Input warping function
  vector STAN_warp_input(vector x, real a){
    return( -1 + 2*inv(1+exp(-a*x)) );
  }

    // Edit a continuous covariate according to sampled uncertainty
  vector STAN_edit_x_cont(
    vector x_cont,
    data array[] int idx_expand,
    data vector teff_obs,
    vector teff)
  {
    int n = num_elements(x_cont);
    vector[n] x_teff_obs = STAN_expand(teff_obs, idx_expand);
    vector[n] x_teff = STAN_expand(teff, idx_expand);
    return(x_cont + x_teff_obs - x_teff);
  }

  // Compute one constant kernel matrix diagonal.
  vector STAN_kernel_const_diag(data array[] int x, data int kernel_type)
  {
    int n = num_elements(x);
    vector[n] K_diag = rep_vector(1.0, n);
    if (kernel_type == 2) {
      vector[n] is_zero;
      for(j in 1:n) {
        is_zero[j] = (x[j]==0);
      }
      K_diag = is_zero; // binary mask: one if both inputs are 0
    }
    return(K_diag);
  }

  // Compute all constant kernel matrices' diagonals.
  array[] vector STAN_kernel_const_all_diag(
    data int n,
    data array[,] int x,
    data array[,] int x_mask,
    data array[,] int components)
  {
    int num_comps = size(components);
    array[num_comps] vector[n] K_const_diag;
    for (j in 1:num_comps) {
      vector[n] K_diag;
      array[9] int opts = components[j];
      int ctype = opts[1];
      int ktype = opts[2];
      int idx_cat = opts[8];
      int idx_cont = opts[9];

      // Compute mask kernel for continuous covariate
      if (idx_cont != 0) {
        K_diag = STAN_kernel_const_diag(x_mask[idx_cont], 2); // 2=binary mask
      } else {
        K_diag = rep_vector(1.0, n);
      }

      // Compute kernel for categorical covariate
      if (ctype == 0 || ctype == 2) {
        K_diag = K_diag .* STAN_kernel_const_diag(x[idx_cat], ktype);
      }
      K_const_diag[j] = K_diag;
    }
    return(K_const_diag);
  }

  // Exponentiated quadratic kernel (with vector inputs)
  vector STAN_kernel_eq_diag(int n, real alpha) {
    return(rep_vector(square(alpha), n));
  }

  // Multiplier matrix to enable variance masking
  vector STAN_kernel_varmask_diag(vector x, real steepness,
    data array[] real vm_params)
  {
    real a = steepness * vm_params[2];
    real r = inv(a)*logit(vm_params[1]);
    int n = num_elements(x);
    vector[n] f_vm = STAN_var_mask(x - r, a);
    return(f_vm .* f_vm);
  }

  // Multiplier matrix to enable heterogeneous effects
  vector STAN_kernel_beta_diag(vector beta, array[] int idx_expand) {
    return(STAN_expand(beta, idx_expand));
  }

  /*
    Compute diagonals of all kernel matrices.
  */
  array[] vector STAN_kernel_all_diag(
    data int n,
    data array[] vector K_const_diag,
    data array[,] int components,
    data array[] vector x,
    data array[] vector x_unnorm,
    array[] real alpha,
    array[] real wrp,
    array[] vector beta,
    array[] vector teff,
    data array[] real vm_params,
    data array[] int idx_expand,
    data array[] vector teff_zero)
  {
    int idx_wrp = 0;
    int idx_alpha = 0;
    int num_comps = size(components);
    array[num_comps] vector[n] KX_diag;

    // Loop through components
    for(j in 1:num_comps){

      // 1. Initialize with constant part of the kernel diagonal
      vector[n] K_diag = K_const_diag[j];
      vector[n] X;

      // 2. Get component properties
      array[9] int opts = components[j];
      int ctype = opts[1];
      int idx_cont = opts[9];
      int is_heter = opts[4];
      int is_warped = opts[5];
      int is_var_masked = opts[6];
      int is_uncrt = opts[7];

      // 3. Pick the possible continuous covariate of this component
      if(ctype != 0){
        if(is_warped){
          X = x_unnorm[idx_cont];
        }else{
          X = x[idx_cont];
        }
      }

      // 4. Handle possible nonstationarity
      if(is_warped){
        real s;
        idx_wrp += 1;

        // 4.1 Handle possible uncertainty
        if(is_uncrt){
          X = STAN_edit_x_cont(X, idx_expand, teff_zero[1], teff[1]);
        }

        // 4.2 Variance masking
        s = wrp[idx_wrp];
        if(is_var_masked){
          K_diag = K_diag .* STAN_kernel_varmask_diag(X, s, vm_params);
        }

        // 4.3 Input warping
        X = STAN_warp_input(X, s);
      }

      // Compute the kernel matrix  diagonal
      idx_alpha += 1;
      if(ctype != 0){
        K_diag = K_diag .* STAN_kernel_eq_diag(n, alpha[idx_alpha]);
      } else {
        K_diag = square(alpha[idx_alpha]) * K_diag;
      }

      // Possible heterogeneity
      if(is_heter){
        K_diag = K_diag .* STAN_kernel_beta_diag(beta[1], idx_expand);
      }

      KX_diag[j] = K_diag; // store kernel matrix diagonal
    }

    return(KX_diag);
  }

  // Categorical zero-sum kernel
  matrix STAN_kernel_zerosum(data array[] int x1, data array[] int x2, data int num_cat) {
    int n1 = size(x1);
    int n2 = size(x2);
    matrix[n1, n2] K;
    for (i in 1:n1) {
      for (j in 1:n2) {
        if (x1[i] == x2[j]) {
          K[i,j] = 1;
        } else {
          K[i,j] = - inv(num_cat - 1);
        }
      }
    }
    return(K);
  }

  // Categorical kernel
  matrix STAN_kernel_cat(data array[] int x1, data array[] int x2) {
    int n1 = size(x1);
    int n2 = size(x2);
    matrix[n1,n2] K;
    for (i in 1:n1) {
      for (j in 1:n2) {
        K[i,j] = (x1[i] == x2[j]);
      }
    }
    return(K);
  }

  // Binary mask kernel
  matrix STAN_kernel_bin(data array[] int x1, data array[] int x2) {
    int n1 = size(x1);
    int n2 = size(x2);
    matrix[n1,n2] K;
    for (i in 1:n1) {
      for (j in 1:n2) {
        K[i,j] = (x1[i] == 0) * (x2[j] == 0);
      }
    }
    return(K);
  }

  // Compute one constant kernel matrix. Does not depend on parameters and
  // therefore this function never needs to be evaluated during sampling.
  matrix STAN_kernel_const(data array[] int x1, data array[] int x2,
    data int kernel_type,  data int ncat)
  {
    int n1 = num_elements(x1);
    int n2 = num_elements(x2);
    matrix[n1, n2] K;
    if (kernel_type == 1) {
      K = STAN_kernel_cat(x1, x2);
    } else if (kernel_type == 2) {
      K = STAN_kernel_bin(x1, x2);
    } else {
      // kernel_type should be 0
      K = STAN_kernel_zerosum(x1, x2, ncat);
    }
    return(K);
  }

  // Compute all constant kernel matrices. These do not depend on parameters and
  // therefore this function never needs to be evaluated during sampling.
  array[] matrix STAN_kernel_const_all(
    data int n1,           data int n2,
    data array[,] int x1,        data array[,] int x2,
    data array[,] int x1_mask,   data array[,] int x2_mask,
    data array[] int num_levels, data array[,] int components)
  {
    int num_comps = size(components);
    array[num_comps] matrix[n1, n2] K_const;
    for (j in 1:num_comps) {
      matrix[n1, n2] K;
      array[9] int opts = components[j];
      int ctype = opts[1];
      int ktype = opts[2];
      int idx_cat = opts[8];
      int idx_cont = opts[9];

      // Compute mask kernel for continuous covariate
      if (idx_cont != 0) {
        K = STAN_kernel_const(x1_mask[idx_cont], x2_mask[idx_cont], 2, 0);
      } else {
        K = rep_matrix(1, n1, n2);
      }

      // Compute kernel for categorical covariate
      if (ctype == 0 || ctype == 2) {
        int M = num_levels[idx_cat];
        K = K .* STAN_kernel_const(x1[idx_cat], x2[idx_cat], ktype, M);
      }
      K_const[j] = K;
    }
    return(K_const);
  }

  // Exponentiated quadratic kernel (with vector inputs)
  matrix STAN_kernel_eq(vector x1, vector x2, real alpha, real ell) {
    return(gp_exp_quad_cov(to_array_1d(x1), to_array_1d(x2), alpha, ell));
  }

  // Multiplier matrix to enable variance masking
  matrix STAN_kernel_varmask(vector x1, vector x2,
    real steepness, data array[] real vm_params)
  {
    real a = steepness * vm_params[2];
    real r = inv(a)*logit(vm_params[1]);
    return(
      to_matrix(to_matrix(STAN_var_mask(x1 - r, a))) *
      transpose(to_matrix(STAN_var_mask(x2 - r, a)))
    );
  }

  // Multiplier matrix to enable heterogeneous effects
  matrix STAN_kernel_beta(vector beta, array[] int idx1_expand, array[] int idx2_expand) {
    return(
      to_matrix(STAN_expand(sqrt(beta), idx1_expand)) *
      transpose(to_matrix(STAN_expand(sqrt(beta), idx2_expand)))
    );
  }

  /*
    Compute all kernel matrices. These depend on parameters and
    therefore this function needs to be evaluated repeatedly during sampling.
  */
  array[] matrix STAN_kernel_all(
    data int n1,
    data int n2,
    data array[] matrix K_const,
    data array[,] int components,
    data array[] vector x1,
    data array[] vector x2,
    data array[] vector x1_unnorm,
    data array[] vector x2_unnorm,
    array[] real alpha,
    array[] real ell,
    array[] real wrp,
    array[] vector beta,
    array[] vector teff,
    data array[] real vm_params,
    data array[] int idx1_expand,
    data array[] int idx2_expand,
    data array[] vector teff_zero)
  {
    int idx_ell = 0;
    int idx_wrp = 0;
    int idx_alpha = 0;
    int num_comps = size(components);
    array[num_comps] matrix[n1, n2] KX;

    // Loop through components
    for(j in 1:num_comps){

      // 1. Initialize with constant part of the kernel matrix
      matrix[n1, n2] K = K_const[j];
      vector[n1] X1;
      vector[n2] X2;

      // 2. Get component properties
      array[9] int opts = components[j];
      int ctype = opts[1];
      int idx_cont = opts[9];
      int is_heter = opts[4];
      int is_warped = opts[5];
      int is_var_masked = opts[6];
      int is_uncrt = opts[7];

      // 3. Pick the possible continuous covariate of this component
      if(ctype != 0){
        if(is_warped){
          X1 = x1_unnorm[idx_cont];
          X2 = x2_unnorm[idx_cont];
        }else{
          X1 = x1[idx_cont];
          X2 = x2[idx_cont];
        }
      }

      // 4. Handle possible nonstationarity
      if(is_warped){
        real s;
        idx_wrp += 1;

        // 4.1 Handle possible uncertainty
        if(is_uncrt){
          X1 = STAN_edit_x_cont(X1, idx1_expand, teff_zero[1], teff[1]);
          X2 = STAN_edit_x_cont(X2, idx2_expand, teff_zero[1], teff[1]);
        }

        // 4.2 Variance masking
        s = wrp[idx_wrp];
        if(is_var_masked){
          K = K .* STAN_kernel_varmask(X1, X2, s, vm_params);
        }

        // 4.3 Input warping
        X1 = STAN_warp_input(X1, s);
        X2 = STAN_warp_input(X2, s);
      }

      // Compute the kernel matrix
      idx_alpha += 1;
      if(ctype != 0){
        idx_ell += 1;
        K = K .* STAN_kernel_eq(X1, X2, alpha[idx_alpha], ell[idx_ell]);
      } else {
        K = square(alpha[idx_alpha]) * K;
      }

      // Possible heterogeneity
      if(is_heter){
        K = K .* STAN_kernel_beta(beta[1], idx1_expand, idx2_expand);
      }

      KX[j] = K; // store kernel matrix
    }

    return(KX);
  }

  // Log prior density to be added to target
  real STAN_log_prior(real x, data array[] int types, data array[] real p) {
    real log_prior = 0;
    real t = x;

    // Possible transform and log of its absolute derivative
    if (types[2]==1){
      log_prior += log(abs(2*x));
      t = square(x);
    }

    // Value of pdf
    if (types[1]==2){
      log_prior += normal_lpdf(t | p[1], p[2]); // 2 = normal
    }else if (types[1]==3){
      log_prior += student_t_lpdf(t | p[1], 0.0, 1.0); // 3 = student-t
    }else if (types[1]==4){
      log_prior += gamma_lpdf(t | p[1], p[2]); // 4 = gamma
    }else if (types[1]==5){
      log_prior += inv_gamma_lpdf(t | p[1], p[2]); // 5 = inv-gamma
    }else if (types[1]==6){
      log_prior += lognormal_lpdf(t | p[1], p[2]); // 6 = log-normal
    }

    return(log_prior);
  }

  // Compute total signal by summing components
  vector STAN_vectorsum(array[] vector vecs, data int L){
    int num_vecs = size(vecs);
    vector[L] s = rep_vector(0, L);
    for (j in 1:num_vecs){
      s += vecs[j];
    }
    return(s);
  }

  // Sum an array of matrices
  matrix STAN_matrix_array_sum(array[] matrix K){
    int n1 = rows(K[1]);
    int n2 = cols(K[1]);
    matrix[n1, n2] K_sum = K[1];
    for(j in 2:size(K)){
      K_sum += K[j];
    }
    return(K_sum);
  }
}
