functions {
  // This function avoids potential numerical problems during warmup 
  // which can appear using the generator poisson_log_rng() that  
  // Underlying idea is to use the Poisson - Gamma mixture
  int neg_binomial_2_log_safe_rng(real eta, real phi) {
    real gamma_rate = gamma_rng(phi, phi / exp(eta));
    if (gamma_rate >= exp(20.79))
      return -9;     
    return poisson_rng(gamma_rate);
  }
}
data {
  int<lower=1> N;                            // Number of data points  
  int<lower=0> complaints[N];                // Outcome variable, number of complaints (vector of integer of size N)         
  vector<lower=0>[N] traps;                  // Number of traps (vector of size N)       
  vector[N] sqfoot;                          // Exposure term
  
  //Building level data
  int<lower=1> B;                           // Number of buildings
  int<lower=1, upper=B> building_idx[N];    // Building identifier 
  int<lower=1> K;                           // Number of columns of the building data matrix
  matrix[B,K] building_data;                // Building data information matrix 
 
  // month info
  int<lower=1> M;                           // Number of months
  int<lower=1,upper=M> mo_idx[N];           // Monthly index
}

transformed data{                           // Taking the log of the square footage
  vector[N] log_sqfoot;  
  log_sqfoot = log(sqfoot);
}

parameters {
  real<lower=0> inv_phi;                    // Inverse of precision parameter (lower-bounded) 
  
  // Varying intercepts parameters
  real alpha;                               // 'global' intercept
  vector[K] zeta;                           // coefficients on building variables in model for mu 
  vector[B] mu_raw;                         // params for non-centered parametrisation of mu
  real<lower=0> sigma_mu;                   // sd of building-specific intercepts        
  
  // Varying slope parameters
  real beta;                                // 'global' slope on traps variable
  vector[K] gamma;                          // coefficients on building-level predictors in model for kappa
  vector[B] kappa_raw;                      // params for non-centered param of building-specific slopes
  real<lower=0> sigma_kappa;                // sd of buildings-specific slopes          
  
  // Monthly effect
  real<lower=0,upper=1> rho_raw;            // used to construct AR(1) coefficient
  vector[M] mo_raw;                         // param for non-centered param of AR(1) process
  real<lower=0> sigma_mo;                   // sd of month-specific parameters
}
transformed parameters {
  real phi = inv(inv_phi);
  
  // NCP of building-specific intercepts and slopes
  vector[B] mu = alpha + building_data * zeta + sigma_mu * mu_raw;
  vector[B] kappa = beta + building_data * gamma + sigma_kappa * kappa_raw;
  
  // AR(1) process priors
  real rho = 2.0 * rho_raw - 1.0;
  vector[M] mo = sigma_mo * mo_raw;
  mo[1] /= sqrt(1 - rho^2);
  for (m in 2:M) {
    mo[m] += rho * mo[m-1];
    // The latter can be equivalently written as
    //mo[m] = mo[m] + rho * mo[m-1];
  }
}

model {
  // Prior for the inverse precision parameter 
  inv_phi ~ normal(0, 1);
  
  // Priors for the varying intercept parameters 
  alpha ~ normal(log(4), 1);
  zeta ~ normal(0, 1);
  mu_raw ~ normal(0,1) ;
  sigma_mu ~ normal(0, 1);
  
  // Priors for the varying slope parameters 
  kappa_raw ~ normal(0,1) ;
  sigma_kappa ~ normal(0, 1);
  beta ~ normal(-0.5, 1);
  gamma ~ normal(0, 1);
  
  // Priors for the monthly effect
  rho_raw ~ beta(10, 5);
  mo_raw ~ normal(0,1);
  sigma_mo ~ normal(0,1);
  
  // Log-likelihood 
  complaints ~ neg_binomial_2_log(mu[building_idx] + kappa[building_idx] .* traps + mo[mo_idx] + log_sqfoot, phi);
}

generated quantities {
  // In-sample replication 
  int y_rep[N];            

  // Log-likelihood contributions                                                              
  vector[N] log_lik;
  
  // Linear predictor (n-th component)
  real eta_n; 

  for (n in 1:N) {
    eta_n = mu[building_idx[n]] + kappa[building_idx[n]] * traps[n] + mo[mo_idx[n]] + log_sqfoot[n];
    y_rep[n] = neg_binomial_2_log_safe_rng(eta_n, phi);
    log_lik[n] = neg_binomial_2_log_lpmf(complaints[n] | eta_n, phi);
  }
}