functions{
  // probabilities in the mosquito feeding cycle model
  vector P_A(vector alpha_B, vector alpha_M, real t){ 
    return exp(-(alpha_B + alpha_M)*t);
  } 
  vector P_B(vector alpha_B, vector alpha_M, real t){
    return (1 - P_A(alpha_B, alpha_M, t)) .* (alpha_B) ./(alpha_B + alpha_M);
  }
  vector P_M(vector alpha_B, vector alpha_M, real t){
    return (1 - P_A(alpha_B, alpha_M, t)) .* (alpha_M) ./(alpha_B + alpha_M);
  }
  // probabilities of the semi-field multinomial model
  vector p_B(vector alpha_B, vector alpha_M){                        
    return P_B(alpha_B, alpha_M, 1);
  }
  vector p_A(vector alpha_B, vector alpha_M){
    return P_A(alpha_B, alpha_M, 1);
  }
  vector p_M(vector alpha_B, vector alpha_M){
    return 1 - p_A(alpha_B, alpha_M) - p_B(alpha_B, alpha_M);
  }
}
data{
int<lower=0> n; // number of experiments per arm
int<lower=0> K; // number observed mosquito endpoints (HLC/fed, resting, knock-down) = 3
int<lower=0> y0[n,K]; // observed endpoints (HLC/fed, resting, knock-down) for n control experiments 
int<lower=0> y1[n,K]; // observed endpoints (HLC/fed, resting, knock-down) for n intervention experiments 
real<lower=0> priorsigma_mean_logrates; // sd of the rate priors
real<lower=0> priorsigma_pikappa; // sd of the pi and kappa priors
real<lower=0> hierarchy; // scale parameter for sigma's
}
parameters{
  real<upper=1> pipi; 
  real<lower=0> kappa;
  real a;
  real b;
  real<lower=0> sigma_a;
  real<lower=0> sigma_b;
  vector[n] phi_a; 
  vector[n] phi_b;
}
transformed parameters{
  vector<lower=0>[n] alpha_B_k0; // rates each arm
  vector<lower=0>[n] alpha_M_k0;
  vector<lower=0>[n] alpha_B_k1;
  vector<lower=0>[n] alpha_M_k1;
  matrix[K,n] theta0; // probabilities of multinomial model each arm and mosquito endpoint
  matrix[K,n] theta1; 

  alpha_B_k0 = exp(a + phi_a *sigma_a); // rates each arm
  alpha_B_k1 = (1 -pipi) *alpha_B_k0;
  alpha_M_k0 = exp(b + phi_b *sigma_b);
  alpha_M_k1 = alpha_M_k0 + kappa * alpha_B_k0;
  theta0[1,] = (p_B(alpha_B_k0, alpha_M_k0))'; // probabilities of multinomial model each arm and mosquito endpoint
  theta0[2,] = (p_A(alpha_B_k0, alpha_M_k0))';
  theta0[3,] = (p_M(alpha_B_k0, alpha_M_k0))';
  theta1[1,] = (p_B(alpha_B_k1, alpha_M_k1))';
  theta1[2,] = (p_A(alpha_B_k1, alpha_M_k1))';
  theta1[3,] = (p_M(alpha_B_k1, alpha_M_k1))';
}
model{
  vector[K] theta_k0; // local variables
  vector[K] theta_k1;
//priors
  target += lognormal_lpdf(1 -pipi | 0, priorsigma_pikappa); // intervention parameters
  target += 2*cauchy_lpdf(kappa | 0,priorsigma_pikappa);
  a ~ normal(0,priorsigma_mean_logrates); // mean of logrates
  b ~ normal(0,priorsigma_mean_logrates);
  phi_a ~ normal(0,1); // normalised deviation from mean logrates
  phi_b ~ normal(0,1);
  target += 2*cauchy_lpdf(sigma_a | 0,hierarchy); // variation of deviation from mean logrates
  target += 2*cauchy_lpdf(sigma_b | 0,hierarchy);

  for (k in 1:n){
    theta_k0 = theta0[,k]; // select data from each arm
    theta_k1 = theta1[,k];
    target += multinomial_lpmf( y0[k,] | theta_k0); // fit multinomial model
    target += multinomial_lpmf( y1[k,] | theta_k1);
  }
}
generated quantities {
  real<lower=0> alpha_B0; // means of rates
  real<lower=0> alpha_M0;
  real<lower=0> alpha_B1;
  real<lower=0> alpha_M1;
  alpha_B0 = exp(a + sigma_a^2/2);
  alpha_M0 = exp(b + sigma_b^2/2);
  alpha_B1 = (1 -pipi) *alpha_B0;
  alpha_M1 = alpha_M0 + kappa * alpha_B0;
}