functions {
  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;                     
  int<lower=0> complaints[N];              
  vector<lower=0>[N] traps;                
  vector[N] sqfoot;  
  
  // building-level data
  int<lower=1> B;
  int<lower=1> K;
  int<lower=1, upper=B> building_idx[N];
  matrix[B,K] building_data;
}

transformed data{
  vector[N] log_sqfoot;  
  log_sqfoot = log(sqfoot);
}

parameters {
 real<lower=0> inv_phi;   

  real alpha;              // 'global' intercept
  vector[K] zeta;          // coefficients on building variables in model for mu 
  vector[B] mu_raw;        // N(0,1) params for non-centered parametrisation of mu
  real<lower=0> sigma_mu;  // sd of building-specific intercepts

  real beta;                 // 'global' slope on traps variable
  vector[K] gamma;           // coefficients on building-level predictors in model for kappa
  vector[B] kappa_raw;       // N(0,1) params for non-centered param of building-specific slopes
  real<lower=0> sigma_kappa; // sd of buildings-specific slopes
}

transformed parameters {
  real phi = inv(inv_phi);
  
  // non-centered parameterization 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;
}

model {
  inv_phi ~ normal(0, 1);
  // mu ~ normal(alpha + building_data * zeta, sigma_mu)
  alpha ~ normal(log(4), 1);
  mu_raw ~ normal(0, 1);   
  sigma_mu ~ normal(0, 1);
  zeta ~ normal(0, 1);

  // kappa ~ normal(beta + building_data * gamma, sigma_kappa)
  beta ~ normal(-0.25, 1);
  kappa_raw ~ normal(0,1) ;
  sigma_kappa ~ normal(0, 1);
  gamma ~ normal(0, 1);
  
  complaints ~ neg_binomial_2_log(mu[building_idx] + kappa[building_idx] .* traps  + log_sqfoot, phi);
} 
generated quantities {
  int y_rep[N];
  vector[N] log_lik;
  for (n in 1:N) {
    real eta_n = mu[building_idx[n]] + kappa[building_idx[n]] * traps[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); 
  }
}