data{
  int N;                                 // number of patients (M(T1))
  int p;                                 // number of covariate (J)
  real<lower=0> t[N];                    // observed time (t^*_{ei} in the paper)
  int<lower=0, upper=1> status[N];       // event indicator (delta_{ei} in the paper: 1 = event, 0 = right censored)
  real<lower=0> a;              // 1st parameter of the gamma prior (shape) for lambda_0e (=alpha_e in the paper)
  real<lower=0> b;              // 2nd parameter of the gamma prior (rate / inverse scale) for lambda_0e (=beta_e in the paper)
  real<lower=0> c;              // 1st parameter of the gamma prior (shape) for rho_e (=alpha_re in the paper)
  real<lower=0> d;              // 2nd parameter of the gamma prior (rate / inverse scale) for rho_e (=beta_re in the paper)
  real mean_c;                  // 1st (mean) parameter of the normal prior for eta
  real<lower=0> se_c;           // 2nd (sd) parameter of the normal prior for eta (=sigma_eta in the paper)
  real mean_b;                  // 1st (mean) parameter of the normal prior for beta_ej
  real<lower=0> se_b;           // 2nd (sd) parameter of the normal prior for beta_ej (= sigma_be)
  real<lower=0,upper=1> p_trt;     // prespecified randomization probability for the experimental arm (= gamma in the paper)
  matrix[N,p] cov;               // covariates (each row is a covariate vector for a patient)
}

parameters{
  real<lower=0> rho;            // shape of weibull: rho_e
  real<lower=0> lambda;     // hazard scale of baseline (lambda_0e)-- warning: this is not the time scale sigma for weibull used in stan
  real<lower=0> eta;                         // treatment effect (eta)
  vector[p] beta;                 // regression coefficients for covariates (beta_e)
}

model{
 rho ~ gamma(c,d);        // rho_e
 lambda ~ gamma(a,b);     // lambda_0e
 eta ~ normal(mean_c,se_c);   // eta
 beta ~ normal(mean_b, se_b);  // beta_e
    for(i in 1:N){
      real linpred_control = dot_product(beta, cov[i]);
      real linpred_exp = eta + linpred_control;
      real hazard_scale_control = lambda*exp(linpred_control);   // the hazard scale for control arm lambda_{ei}(x_i,Z_i) as in equation (2) in the paper
      real time_scale_control   = pow(hazard_scale_control, -1.0 / rho);     // (hazard_scale)^(-1/rho)
      real hazard_scale_exp = lambda*exp(linpred_exp);   // the hazard scale for exp arm lambda_{ei}(x_i,Z_i) as in equation (2) in the paper
      real time_scale_exp   = pow(hazard_scale_exp, -1.0 / rho);     // (hazard_scale)^(-1/rho)
    if(status[i]==1){
      target += log_mix(p_trt,
                     weibull_lpdf(t[i] | rho, time_scale_exp),
                     weibull_lpdf(t[i] | rho, time_scale_control));
    } else {
      target += log_mix(p_trt,
                     weibull_lccdf(t[i] | rho, time_scale_exp),
                     weibull_lccdf(t[i] | rho, time_scale_control));
    }
    }
}

generated quantities { // unknown treatment assignment (x_i)
  int<lower=0,upper=1> x[N];   // unknown treatment assignment (x_i)
  real<lower=0,upper=1> p_x1[N];   // Pr(x_i=1|rest parameters, data) for this draw
  //real<lower=0,upper=1> p_x1_test[N];   // Pr(x_i=1|rest parameters, data) for this draw
  for(i in 1:N){
    real linpred_control = dot_product(beta, cov[i]);
    real linpred_exp = eta + linpred_control;
    real hazard_scale_control = lambda*exp(linpred_control);   // the hazard scale for control arm lambda_{ei}(x_i,Z_i) as in equation (2) in the paper
    real time_scale_control   = pow(hazard_scale_control, -1.0 / rho);     // (hazard_scale)^(-1/rho)
    real hazard_scale_exp = lambda*exp(linpred_exp);   // the hazard scale for exp arm lambda_{ei}(x_i,Z_i) as in equation (2) in the paper
    real time_scale_exp   = pow(hazard_scale_exp, -1.0 / rho);     // (hazard_scale)^(-1/rho)

    // log likelihood under each arm for this subject (conditional on x)
    real l1;
    real l0;
    if (status[i] == 1) {
      l1 = weibull_lpdf(t[i] | rho, time_scale_exp);
      l0 = weibull_lpdf(t[i] | rho, time_scale_control);
    } else {
      l1 = weibull_lccdf(t[i] | rho, time_scale_exp);
      l0 = weibull_lccdf(t[i] | rho, time_scale_control);
    }

    // posterior Pr(x_i=1 | rest)
    //p_x1_test[i] = p_trt*exp(l1)/(p_trt*exp(l1)+(1-p_trt)*exp(l0));
    p_x1[i] = inv_logit( logit(p_trt) + (l1 - l0) );

    // draw a discrete assignment for downstream prediction
    x[i] = bernoulli_rng(p_x1[i]);

  }
}
