// -*- mode: C++; -*-
/* ############################################################
 * Random slopes model with Brownian motion errors
 * ###########################################################
 * Assume that there is a random slope and that the initial time
 * is common accross all students.
 * Assume that the latent variable has mean 0 and std 1 at Time 1.
 * This version has the EM parameter 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 */
  real Obs[I,Tmax];           // Observed scores
  real Time[I,Tmax-1];          // Time between measures
  real Dose[I,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 T0;                      // Effective start time.
  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 {
  real T0k;                     // Offset constant to make starting
                                // ability have mean 0.

  /* 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.

  /* Latent model derived constants. */
  T0k = slope_mu_mu*T0;

  /* 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 treat_eff;               // Treatment effect
  real<lower=0.0,upper=1.0> T0_err_std;     // Error std at time 0.
  real slope_mu;                // Mean slope
  /* Brownian Motion Level 1 parameters */
  real slope[I];                // Slope varies by person
  /* Brownian Process latent variables. */
  real innov[I,Tmax-1];           // Innovations for each time step.
  real T0_err[I];              // Error at time 0.



}
transformed parameters {
  /* Constrain T0^2*var(slope) + var(T0_err) = 1 */
  real<lower=0.0> slope_std;    // STD of slope.
  /* Latent variable for Markov process. */
  real theta[I,Tmax];           // Latent variable.


  /* Constrain T0^2*var(slope) + var(T0_err) = 1 */
  slope_std = sqrt(1-T0_err_std^2)/T0;

  /* Latent variable for Markov process. */
  for (i in 1:I) {
    theta[i,1] = slope[i]*T0-T0k+T0_err[i];
    for (t in 2:T[i]) {
      theta[i,t] = theta[i,t-1]
                   + slope[i]*Time[i,t-1]
                   + treat_eff*Dose[i,t-1]
                   + innov[i,t-1];
    }
  }
}
model {

  /* Higher level parameters */
  slope_mu ~ normal(slope_mu_mu,slope_mu_std);
  
  /* Parameter Priors */
  slope ~ normal(slope_mu,slope_std);
    
  /* Innovations */
  T0_err ~ normal(0.0,T0_err_std);
  for (i in 1:I) {
    for (t in 1:(T[i]-1)) {
      innov[i,t] ~ normal(0.0,sqrt(Time[i,t]*var_innov));
    }
  }

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

}
generated quantities{

}