// Online resource 4
// 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 M1. This Stan code is used in the mixed model.

data {
int<lower=1> P; // number of fixed effects
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[P] x[N]; // N * P design matrix for fixed effects
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 {
matrix[K, P] coef; // K * P matrix of fixed-effect coefficients (to be estimated)  

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
matrix[K, K] Sigma; 
vector[K] mu[N]; // N * K matrix of predicted means



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]] + coef * x[n];
	}  

	 // define the observation level covariance here for use in the multivariate T
  Sigma = tcrossprod(L_Sigma); 
}

model {
// Variance components for residuals. 
// For lkj: parameter=1 gives uniform prior, 
// parameter=2 makes correlations off diagonal close to zero.
L_Omega ~ lkj_corr_cholesky(1.5); 
L_sigma ~ cauchy(0, 2.5);

// Specimen variance components.
L_Phi ~ lkj_corr_cholesky(1.5);
L_theta ~ cauchy(0, 2.5);

// Lot variance components.
L_Zeta ~ lkj_corr_cholesky(1.5);
L_beta ~ cauchy(0, 2.5);

mu_specimen ~ multi_normal_cholesky(zeros, L_Theta);
mu_lot ~ multi_normal_cholesky(zeros, L_Beta);
to_vector(coef) ~ normal(0, 5); 

// 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; 
vector[K] resids[N]; // N * K matrix of raw residuals
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]);   
	ICC_lot[k] = Beta[k, k] / (Vcov_obs[k, k] + Beta[k, k]+ Theta[k, k]);
	}

for (n in 1:N) {
	resids[n] = y[n] - mu[n];
	Maha_sqd[n] = quad_form_sym(Sigma_inv, resids[n,]); 
	log_lik[n] = multi_student_t_lpdf(y[n] | DF_obs, mu[n], Sigma); 
	}
}