model{
  
  # prior distributions
  alpha ~ dnorm(0.05, 1536.292) T(0, )
  for(k in 1:K){
    b[k] ~ dlogis(0, 1)
  }
  
  ### CALCULATING STABLE STAGE DISTRIBUTION
  for(i in 1:G){  # looping over all groups
    
    lambda[i] ~ dnorm(1, 1536.292)
    
    # survival probability, 1 time step
    logit(p[i, 1]) <- b[1] + b[5] * mmr[i] + b[6] * umr[i]
    
    # survive to next stage class
    g[i, 1] <- p[i, 1] ^ d[1] * (1 - p[i, 1]) / (1 - p[i, 1] ^ d[1])
    
    # remain in current stage class
    z[i, 1] <- (1 - p[i, 1] ^ (d[1] - 1)) / (1 - p[i, 1] ^ d[1]) * p[i, 1]
    
    w[i, 1] <- 1  # proportion to stable stage
    
    for(s in 2:(S - 1)){  # loop through all stage classes
      logit(p[i, s]) <- b[1] + b[s] + b[5] * mmr[i] + b[6] * umr[i]
      g[i, s] <- p[i, s] ^ d[s] * (1 - p[i, s]) / (1 - p[i, s] ^ d[s])
      z[i, s] <- (1 - p[i, s] ^ (d[s] - 1)) / (1 - p[i, s] ^ d[s]) * p[i, s]
      w[i, s] <- g[i, s - 1] / (lambda[i] - z[i, s]) * w[i, s - 1]
    }
    
    logit(p[i, S]) <- b[1] + b[S - 1] + b[5] * mmr[i] + b[6] * umr[i]
    g[i, S] <- 0
    z[i, S] <- p[i, S]
    w[i, S] <- g[i, S - 1] / (lambda[i] - z[i, S]) * w[i, S - 1]
    
    # stable stage distribution
    C[i, 1:S] <- w[i, 1:S] / sum(w[i, 1:S])
    
    ### EXPECTED COUNT IN EACH STAGE CLASS
    ey[i, 2] <- g[i, 1] * C[i, 1] + z[i, 2] * C[i, 2]
    ey[i, 3] <- g[i, 2] * C[i, 2] + z[i, 3] * C[i, 3]
    ey[i, 4] <- g[i, 3] * C[i, 3] + z[i, 4] * C[i, 4]
    ey[i, 5] <- g[i, 4] * C[i, 3] + z[i, 5] * C[i, 5]
    ey[i, 1] <- alpha * (lambda[i] - sum(ey[i, 2:5]))
    
    ### LIKELIHOOD
    theta[i, 1:S] ~ ddirch(ey[i, 1:S])
    y[i, 1:S] ~ dmulti(theta[i, 1:S], N[i])
    
    ### MODEL CHECKING
    y_new[i, 1:S] ~ dmulti(theta[i, 1:S], N[i])
    
    for(n in 1:S){
    
      g_obs[i, n] <- y[i, n] *
        log(y[i, n] / (N[i] * ey[i, n] / sum(ey[i, 1:S])))
      
      g_new[i, n] <- y_new[i, n] *
        log(y_new[i, n] / (N[i] * ey[i, n] / sum(ey[i, 1:S])))
      
    }
    
  }
  
  cs_obs <- 2 * sum(g_obs)
  cs_new <- 2 * sum(g_new)
  bp <- step(cs_new - cs_obs)
  
}