// Online resource 3
// Title: A method for the taphonomic assessment of bone tools using 3D surface texture analysis of bone microtopography
// Authors: Naomi L. Martisius a,b*, Shannon P. McPherron b, Ellen Schulz-Kornas c,d,b, Marie Soressi e,b, Teresa E. Steele a,b
// Affiliations: a Department of Anthropology, University of California, Davis, Davis, CA, USA; b Department of Human Evolution, Max Planck Institute for Evolutionary Anthropology, Leipzig, Germany; c Department of Cariology, Endodontology and Periodontology, University of Leipzig, Leipzig, Germany; d Max Planck Weizmann Center for Integrative Archaeology and Anthropology, Max Planck Institute for Evolutionary Anthropology, Leipzig, Germany; e Faculty of Archaeology, Leiden University, Leiden, The Netherlands
// * Corresponding author: Naomi L. Martisius (nlmartisius@ucdavis.edu)

// Stan code M0. This Stan code is used in the empty model. 

data {
  int<lower=1> K; // number of features
  int<lower=1> J; // number of specimens (for random intercepts)
  int<lower=1> M; // number of lots (for random intercepts)
  int<lower=0> N; // number of observations (sample size) 
  vector[K] y[N]; // N * K matrix of observations
  int<lower=1> specimen[N]; // vector of integer-coded specimen IDs for each observation
  int<lower=1> lot[N]; // vector of integer-coded lots for each observation
}

transformed data {
  vector[K] zeros;
  zeros = rep_vector(0, K);
}

parameters {
  cholesky_factor_corr[K] L_Omega; // for residual vectors
  vector<lower=0>[K] L_sigma; // for residual vectors
  
  cholesky_factor_corr[K] L_Phi; // for specimen random intercepts
  vector<lower=0>[K] L_theta; // for specimen random intercepts
  vector[K] mu_specimen[J]; // J * K matrix of specimen random intercepts
  
  cholesky_factor_corr[K] L_Zeta; // for lot random intercepts
  vector<lower=0>[K] L_beta; // for lot random intercepts
  vector[K] mu_lot[M]; // M * K matrix of lot random intercepts
  
  real<lower=2>nu_obs; //degrees of freedom parameter for multivariate T
}


transformed parameters {
  // Definitions and calculations are here, rather than in the model block, so that 
  // these are available for generated quantities. 
  matrix[K, K] L_Sigma; // for residual vectors
  matrix[K, K] L_Theta; // for specimen random intercepts
  matrix[K, K] L_Beta; // for lot random intercepts
  vector[K] mu[N]; // N * K matrix of predicted means
  vector[K] resids[N]; // N * K matrix of raw residuals

  matrix[K, K] Sigma; 
  
  L_Sigma = diag_pre_multiply(L_sigma, L_Omega);
  L_Theta = diag_pre_multiply(L_theta, L_Phi);
  L_Beta = diag_pre_multiply(L_beta, L_Zeta);
  for (n in 1:N){
    // Predicted mean contains nested reference for the specimen intercept 
    // along with nested reference for lot intercept. 
    mu[n] = mu_specimen[specimen[n]] + mu_lot[lot[n]];
    resids[n] = y[n] - mu[n];
  }  
  
  // define the observation level covariance here for use in the multivariate T
  Sigma = tcrossprod(L_Sigma); 
  
}

model {
  // Nested reference for the specimen intercept + nested reference for lot 
  
  // Variance components for residuals.
  L_Omega ~ lkj_corr_cholesky(1.5); // for lkj prior=1 means uniform prior, if prior=2 then correlations off diagonal close to zero
  L_sigma ~ cauchy(0, 25);
  
  // Specimen variance components.
  L_Phi ~ lkj_corr_cholesky(1.5);
  L_theta ~ cauchy(0, 25);
  
  // Lot variance components.
  L_Zeta ~ lkj_corr_cholesky(1.5);
  L_beta ~ cauchy(0, 25);
  
  mu_specimen ~ multi_normal_cholesky(zeros, L_Theta);
  mu_lot ~ multi_normal_cholesky(zeros, L_Beta);
  // Multivariate T. 
  nu_obs ~ gamma(2, 0.1); // recommended by Aki on Gelman's blog, 17 May 2015
  
  y ~ multi_student_t(nu_obs, mu, Sigma); // Model of observations
}

generated quantities {
  matrix[K, K] Vcov_obs;
  matrix[K, K] Theta;
  matrix[K, K] Beta;
  vector[K] ICC_specimen; 
  vector[K] ICC_lot;
  matrix[K, K] L_Sigma_inv; // Inverse Cholesky factor of Sigma.
  matrix[K, K] Sigma_inv;
  // Squared Mahalanobis distance of each observation from its predicted value.  
  vector[N] Maha_sqd;
  vector[N] log_lik; 
  real DF_obs; // degrees of freedom for observation level multivariate T
  real IF_obs; // inflation factor for the scale matrix
  
  Theta = tcrossprod(L_Theta);
  Beta = tcrossprod(L_Beta);
  L_Sigma_inv = inverse(L_Sigma);
  Sigma_inv = crossprod(L_Sigma_inv); // Inverse covariance matrix from inverted Cholesky factor.
  DF_obs = nu_obs; 
  IF_obs = DF_obs/(DF_obs-2);
  Vcov_obs = IF_obs*Sigma; 
  for (k in 1:K) ICC_specimen[k] = Theta[k, k] / (Vcov_obs[k, k] + Beta[k, k]+ Theta[k, k]);   
  for (k in 1:K) ICC_lot[k] = Beta[k, k] / (Vcov_obs[k, k] + Beta[k, k]+ Theta[k, k]);
  for (n in 1:N) Maha_sqd[n] = quad_form_sym(Sigma_inv, resids[n,]); 
  for (n in 1:N) log_lik[n] = multi_student_t_lpdf(y[n] | DF_obs, mu[n], Sigma); 
}