data {
  int<lower=0> R;       // number of regions
  int<lower=0> S;       // number of states (either 1 or 27)
  int<lower=0> A;       // number of ages (usually 100)
  int<lower=0> K;       // number of basis functions (usually 7)
  
  matrix[A,K] B;             // spline basis functions (usually 100 x 7)
  vector[A] std_schedule;    // std logmu schedule for TOPALS

  vector[R] agg_wts;       // wts for aggregating pi[3]s up to state level (won't be used if state is the 'small' geog)
  vector[3] death_frac;    // fraction of deaths at ages 0, 1-29, 30+
  
  vector<lower=0,upper=1>[R] pihat1;    //  pi[r,1] ~ beta( K1*pihat1[r], K1*(1-pihat1[r]))
  vector<lower=0,upper=1>[R] pihat_all;  //  pi[r,ALL] ~ beta( Kall*pihat_all[r], Kall*(1-pihat_all[r]))
  
  vector[S] a3 ;                      //  pibar[3] ~ beta(a3,b3)  state agg
  vector[S] b3 ;          
  
  matrix<lower=0>[A,R] N      ;  // age- and region-specific exposure
  int<lower=0>         D[A,R] ;  // observed deaths by age and region

}

transformed data {
  matrix[A,R] lambda_star;   // std logmu schedule for TOPALS in each column
  
  lambda_star = rep_matrix(std_schedule, R);
}

parameters {
  matrix[K,R] alpha;   // TOPALS offsets
  
  simplex[4] eta[R]; // used to construct (pi1, pi2, pi3)
  
  real<lower=0> K1;    // precision (hyper)parameter for regional pi1 
  real<lower=0> Kall;  // precision (hyper)parameter for regional piall
}

transformed parameters {
  matrix<lower=0,upper=1>[R,3] pi; 
  vector<lower=0,upper=1>[R]   piall;    // coverage probabilities at all ages combined
  real<lower=0, upper=1>       agg_pi3;  // state-level coverage for deaths at 30+
  
  for (region in 1:R) {
  
    pi[region,1] = eta[region,1];
    pi[region,2] = eta[region,1] + eta[region,2] + eta[region,3];  
    pi[region,3] = eta[region,1] + eta[region,2];    
  
  } # for region
  
  piall = pi * death_frac;   // result is an R-vector of avg covg (all ages) by region
  
  agg_pi3 = dot_product( pi[,3], agg_wts); 
  
}

model {
  matrix[A,R]    lambda;               // log mx rates
  matrix[K-1,R]  alpha_diff;           // differences in offsets over adjacent groups
  
  lambda      = lambda_star + B * alpha ;   // regional logmu schedules ages 0,1,..99

  alpha_diff  = alpha[2:K,] - alpha[1:(K-1),];   // first diffs in alphas
  
  ## LIKELIHOOD: Poisson distribution of deaths    
  ## log(N) + lambda + log(pi) = ln(expect death reports)       
  
  for (region in 1:R) {
    D[1,region]    ~ poisson_log( log(N[1,region])    + lambda[1,region]     + log(pi[region,1])) ;  
    D[2:30,region] ~ poisson_log( log(N[2:30,region]) + lambda[2:30,region]  + log(pi[region,2])) ;  
    D[31:A,region] ~ poisson_log( log(N[31:A,region]) + lambda[31:A, region] + log(pi[region,3])) ;  
  }
  
  ## PRIORS
  to_vector(alpha)       ~ normal(0,4);
  to_vector(alpha_diff)  ~ normal(0, 1/sqrt(2));

  K1      ~ exponential(0.05) ;   // truncated exponential for precision of K1
  Kall    ~ exponential(0.05) ;   // truncated exponential for precision of Kall
  
  pi[,1] ~ beta(     (5+K1)*pihat1, (5+K1)*(1-pihat1) );
  
  piall  ~ beta( (5+Kall)*pihat_all, (5+Kall)*(1-pihat_all)  ) ;    // vectorized: prior for total covg regions 1..R
  
  if (S  > 1) pi[,3]  ~ beta(a3, b3 );    // prior for state-level pi3 values (if state is the SMALL area)
  if (S == 1) agg_pi3 ~ beta( a3, b3) ;   // prior for state-level pi3 value  (if state is the LARGE area)
  
  

} //model