data {
  int<lower=1> K_loc;
  int<lower=1> K_scale;
  int<lower=1> K_total;
  int<lower=1> N; //total number of observations
  int<lower=1> nsubj; //number of subjects
  int<lower=0> n[nsubj + 1]; 
  //cumulative sums of observations per subject, 
  //with a 0 at the beginning for counting purposes
  int<lower=1, upper = nsubj> subject[N];
  int<lower=0> time[N];
  real<lower=0> y[N];
}

parameters {
  //location
  vector[K_loc] beta0;
  vector[K_loc] beta1;
  
  //scale
  vector[K_scale] tau0;
  vector[K_scale] tau1;
  
  //mixing proportions
  simplex[K_loc] pi_class_location;
  simplex[K_scale] pi_class_scale;
}

model{
  vector[K_loc] log_pi_location = log(pi_class_location); 
  vector[K_scale] log_pi_scale = log(pi_class_scale);//cache log calculation
  vector[K_total] log_pi;
  
  //likelihood
  for (j in 1:nsubj){
    vector[K_total] lps = rep_vector(0, K_total);
    
    for(i in (n[j]+1):n[j + 1]){
      for(k_loc in 1:K_loc){
        for(k_scale in 1:K_scale){
          log_pi[(k_loc - 1)*K_scale + k_scale] = log_pi_location[k_loc] 
          + log_pi_scale[k_scale];
          
          lps[(k_loc - 1)*K_scale + k_scale] +=  normal_lpdf(y[i] | 
          beta0[k_loc] + beta1[k_loc] * time[i], 
          sqrt(exp(tau0[k_scale] + tau1[k_scale] * time[i])));
        }
      }
    }
    
    target += log_sum_exp(log_pi + lps);
  }
}	

generated quantities{
  vector[nsubj] log_lik;
  vector[K_loc] log_pi_location = log(pi_class_location);
  vector[K_scale] log_pi_scale = log(pi_class_scale);//cache log calculation
  vector[K_total] log_pi;
  vector[K_loc] pi_class_location_subj[nsubj]; 
  vector[K_scale] pi_class_scale_subj[nsubj]; 
  //posterior classifying probabilities for each subject
  matrix[K_scale, K_loc] scale_pi;
  matrix[K_scale, K_loc] scale_lps;
  for(k_loc in 1:K_loc){
    for(k_scale in 1:K_scale){
      log_pi[(k_loc - 1)*K_scale + k_scale] = log_pi_location[k_loc] + 
      log_pi_scale[k_scale];
      scale_pi[k_scale, k_loc] = log_pi_location[k_loc] + 
      log_pi_scale[k_scale];
    }
  }
  
  for (j in 1:nsubj){
    vector[K_total] lps = rep_vector(0, K_total);
    
    for(k_loc in 1:K_loc){
      for(k_scale in 1:K_scale){
      scale_lps[k_scale, k_loc] = 0;
      }
    }
    
    for(i in (n[j]+1):n[j + 1]){
      for(k_loc in 1:K_loc){
        for(k_scale in 1:K_scale){
          scale_lps[k_scale, k_loc] += normal_lpdf(y[i] | 
          beta0[k_loc] + beta1[k_loc] * time[i],
          sqrt(exp(tau0[k_scale] + tau1[k_scale] * time[i])));
          
          lps[(k_loc - 1)*K_scale + k_scale] +=  normal_lpdf(y[i] | 
          beta0[k_loc] + beta1[k_loc] * time[i],
          sqrt(exp(tau0[k_scale] + tau1[k_scale] * time[i])));
        }
      }
    }
    
    log_lik[j] = log_sum_exp(log_pi + lps);
    
    for(k_loc in 1:K_loc){
      pi_class_location_subj[j,k_loc] = 
      exp(log_sum_exp(log_pi[((k_loc -1) * 
      K_scale + 1):(k_loc * K_scale)] + 
      lps[((k_loc -1) * K_scale + 1):(k_loc * K_scale)]) - log_lik[j]);
    }
    
    for(k_scale in 1:K_scale){
      pi_class_scale_subj[j,k_scale] = 
      exp(log_sum_exp(scale_pi[k_scale, 1:K_loc] + 
      scale_lps[k_scale, 1:K_loc]) - log_lik[j]);
    }
  }
}