/**********************************************************************
 * linear mixed-effect model 
 * Peng Zeng @ Auburn University 
 * 2025-09-12 
 * 
 * model: 
 *     y[i] ~ normal(b[ID[i]] * x, tau)
 *     b[i] ~ normal(theta, Sigma)
 * prior distribution: 
 *     theta ~ normal(mu0, Lambda0)
 *     Sigma ~ inverse Wishart(eta0, S0)
 *     tau^2 ~ inverse chi-square(v0, sigma0^2)
 **********************************************************************/

data {
    int<lower=0> n;                    // number of observations 
    int<lower=0> p;                    // number of predictors 
    int<lower=0> m;                    // number of groups 
    vector[n] y;                       // response 
    matrix[n, p] x;                    // predictor  
    int<lower=1, upper=m> ID[n];       // ID in {1, ..., m} 
    vector[p] theta_mean;              // mu0 
    cov_matrix[p] theta_var;           // Lambda0 
    real Sigma_df;                     // eta0                        
    cov_matrix[p] Sigma_scale;         // S0 
    real tau2_df;                      // v0  
    real<lower=0> tau2_scale;          // sigma0  
}
parameters {
    array[m] vector[p] beta; 
    cov_matrix[p] Sigma; 
    real<lower=0> tau; 
    vector[p] theta; 
}
model {
    tau^2 ~ scaled_inv_chi_square(tau2_df, tau2_scale); 
    theta ~ multi_normal(theta_mean, theta_var);
    Sigma ~ inv_wishart(Sigma_df, Sigma_scale); 

    beta ~ multi_normal(theta, Sigma); 

    for(i in 1:n) 
    {
        y[i] ~ normal(x[i] * beta[ID[i]], tau); 
    }
}

/**********************************************************************
 * THE END 
 **********************************************************************/