
data { 
  int<lower=0> N;   // number of observations
  int<lower=1> p;   // number of predictors
  int<lower=1> J;   // number of groups in data (persons)
  int<lower=1> k;   // number of group-level predictors
  int<lower=1,upper=J> group[N]; //group indicator
  matrix[N,p] X;    // fixed-effect design matrix
  matrix[N,k] Z;    // random-effect design matrix
  vector[N] Y;      // response
  real<lower=0> scale_icept;    // prior std for the intercept
  real<lower=0> scale_global;   // scale for the half -t prior for tau
  real<lower=1> nu_global;      // degrees of freedom for the half -t priors for tau
  real<lower=1> nu_local;       // degrees of freedom for the half - t priors for lambdas
  real<lower=0> slab_scale;     // slab scale for the regularized horseshoe
  real<lower=0> slab_df;        // slab degrees of freedom for the regularized horseshoe
} 

transformed data { 
  int Pc; 
  matrix[N, p - 1] Xc;    
  vector[p - 1] means_X;  // column means of X before centering 
  
  Pc = p - 1;  // the intercept is removed from the design matrix 
  for (i in 2:p) { 
     means_X[i - 1] = mean(X[, i]); 
     Xc[, i - 1] = X[, i] - means_X[i - 1]; 
  } 
}

parameters { 

  // horseshoe shrinkage parameters 
  real <lower=0> tau;             // global shrinkage parameter
  vector <lower=0>[Pc] lambda;    // local shrinkage parameter
  real <lower=0> caux;

  real temp_Intercept;            // temporary intercept 
  cholesky_factor_corr[k] L;      // Cholesky factor of group ranef corr matrix
  vector<lower=0>[k] sigma_b;     // group-level random-effect standard deviations
  real<lower=0> logsigma_e;       // residual standard deviations 
  vector[Pc] zbeta;                 
  vector[k] z[J];                 // unscaled group-level effects
}

transformed parameters {

  // Latent variables

  vector[Pc] beta;                // population-level effects (fixed effects)
  vector[k] b[J];                 // group-level effects (random effects)

  matrix[k, k] Sigma_b;           // Variance-covariance matrix of group-level effects
  real<lower=0> sigma_e;          // residual standard deviations 

  // regularized horseshoe parameters
  real<lower=0> c;                // slab scale 
  vector<lower=0>[Pc] lambda_tilde;        // ' truncated ' local shrinkage parameter 

  sigma_e = exp(logsigma_e);
  
  // Apply regularizing horseshoe prior for betas
  c = slab_scale * sqrt(caux);
  lambda_tilde = sqrt(c^2 * square(lambda) ./ (c^2 + tau^2* square(lambda)));
  beta = zbeta .* lambda_tilde * tau;

  // Premultiply diagonal matrix [sigma_b] with the Cholesky decomposition L of
  // the correlation matrix Sigma_b to get variance-covariance matrix of group-level effects

  // diag(Sigma_b) * L
  Sigma_b = diag_pre_multiply(sigma_b, L); 
  
  // Group-level effects are generated by multipying D (Sigma_b) with z 
  // that has standard normal distribution
    
  for(j in 1:J) 
    b[j] = Sigma_b * z[j];    
}

model { 
  //vector[N] e; 
  vector[N] mu;                   

  // Horseshoe
  // half-t priors for lambdas and tau, and inverse-gamma for c^2
  zbeta ~ normal(0, 1); 
  lambda ~ student_t(nu_local, 0, 1);
  tau ~ student_t(nu_global, 0, scale_global * sigma_e);
  caux ~ inv_gamma(0.5 * slab_df, 0.5 * slab_df );

  sigma_b ~ student_t(3, 0, 10);
  L ~ lkj_corr_cholesky(1); 
  
  // Standard normal prior for random effects
  for (j in 1:J)
    z[j] ~ normal(0,1);

  // Likelihood 

  // - link function (identity function) for typical correlation
  mu = temp_Intercept + Xc * beta;

  // - add group effects
  for (i in 1:N) 
  {
     mu[i] = mu[i] + Z[i] * b[group[i]];
     
     // residuals
     //e[i] = Y[i] - mu[i];
  }

  // Y and mu are vectors, sigma_e is a scalar that is estimated for whole vector
  Y ~ normal(mu, sigma_e);
}

generated quantities { 
  real beta_Intercept;            // Population-level intercept 
  corr_matrix[k] C; 
  
  // Correlation matrix of random-effects, C = L'L
  C = multiply_lower_tri_self_transpose(L); 
  
  beta_Intercept = temp_Intercept - dot_product(means_X, beta);
} 
