// -*- mode: C++; -*-
/* ############################################################
 * Random slopes model with constant errors
 * ###########################################################
 * Ignore time 0, and just model correlation between slope and
 * intercept. 
 * Assume that the latent variable has mean 0 and std 1 at Time 1.
 * This version has the EM parameters fixed
 */
functions {
}
data {
  int<lower=1> I;               // number of students
  int<lower=1> T[I];            // number of measures per student
  int<lower=1> Tmax;            // Tmax <- max(T);

  /* Data */
  vector[I] Obs[Tmax];           // Observed scores
  vector[I] Time[Tmax-1];          // Time between measures
  vector[I] Dose[Tmax-1];         // Treatment exposure between time points
  /* Note that this is a ragged array, so that Y[i,t] for t>T[i] will be NaN */

  /* Hyperparameters */
  real slope_mu_mu;             // Slope prior mean
  real slope_mu_std;            // Slope prior std.

  /* Evidence model characterisitcs */
  real<lower=0.0,upper=1.0> obs_rel; // Reliability
  real obs_mean_t1;                  // Mean of observations at time 1
  real<lower=0.0> obs_std_t1;    // Standard deviation of observations
                                 // at time 1; 
  
}

transformed data {
  /* Cumulative exposures */
  vector[I] CumTime[Tmax];          // Time between measures
  vector[I] CumDose[Tmax];         // Treatment exposure between time points
  /* Evidence model parameters (linear evidence model)*/
  real<lower=0.0> res_std;      // Residual Standard deviation.
  real obs_slope;               // Slope between observations and
                                // latent variable.
  real obs_int;                 // Intercept of relationshiop between
                                // observations and latent varibles.
  /* Cumulative dosages */
  for (i in 1:I) {
    CumTime[1,i] = 0;
    CumDose[1,i] = 0;
    for (t in 2:T[i]) {
      CumTime[t,i] = Time[t-1,i]+CumTime[t-1,i];
      CumDose[t,i] = Dose[t-1,i]+CumDose[t-1,i];
    }
  }

  /* Evidence model parameters (linear evidence model)*/
  res_std = obs_std_t1*sqrt(1-obs_rel);
  obs_slope = obs_std_t1*sqrt(obs_rel);
  obs_int = obs_mean_t1;
}
parameters {

  /* Brownian Motion Process Level 2 Parameters */
  real<lower=0.0> var_innov;     // Variance of the innovations
  real<lower=0.0> treat_eff;               // Treatment effect
  real slope_mu;                // Mean slope
  real<lower=0.0> slope_std;    // STD of slope.
  real<lower=0.0,upper=1.0> slope_r2; // Squared correlation of slope
                                      // with intercept 
  /* Level 1 growth parameters */
  vector[I] slope_res;                // Slope varies by person
  vector[I] T0_val;              // Value of theta at time 0.
  /* Level 1 residuals latent variables. */
  vector[I] innov[Tmax-1];           // Innovations for each time step.

}
transformed parameters {
  /* Constrain T0^2*var(slope) + var(T0_val) = 1 */
  /* Latent variable for Markov process. */
  vector[I] theta[Tmax];           // Latent variable.
  /* Brownian Motion Level 1 parameters */
  vector[I] slope;                // Slope varies by person


  /* Constrain T0^2*var(slope) + var(T0_val) = 1 */
  slope = slope_mu
    + slope_std*(sqrt(1-slope_r2)*slope_res
                      + slope_r2*T0_val);

  /* Latent variable for Markov process. */
  for (i in 1:I) {
    theta[1,i] = T0_val[i];
    for (t in 2:T[i]) {
    theta[t,i] = T0_val[i] +
        + slope[i]*CumTime[t,i]
        + treat_eff*CumDose[t,i]
        + innov[t-1,i]*sqrt(var_innov);
    }
  }
}
model {

  /* Higher level parameters */
  slope_mu ~ normal(slope_mu_mu,slope_mu_std);
  
  /* Parameter Priors */
  slope_res ~ normal(0.0,1.0);
    
  /* Innovations */
  T0_val ~ normal(0.0,1.0);
  for (i in 1:I) {
    for (t in 1:(T[i]-1)) {
      innov[t,i] ~ normal(0.0,1.0);
    }
  }

  /* Observations */
  for (i in 1:I) {
    for (t in 1:T[i]) {
      Obs[t,i] ~ normal(obs_int+obs_slope*theta[t,i],res_std);
    }
  }

}
generated quantities{

}