//
// This Stan program defines a simple model, with a
// vector of values 'y' modeled as normally distributed
// with mean 'mu' and standard deviation 'sigma'.
//
// Learn more about model development with Stan at:
//
//    http://mc-stan.org/users/interfaces/rstan.html
//    https://github.com/stan-dev/rstan/wiki/RStan-Getting-Started
//

// The input data is a vector 'y' of length 'N'.
data {
  int<lower=1> D; // number of features
  int<lower=1> N; // number of observations
  int<lower=1> L; // number of categories
  array[N] int<lower=0, upper=1> y; // vector of dependent variables
  array[N] int<lower=1, upper=L> ll; // vector of categories
  array[N] row_vector[D] x; // independent variables
  real mu_mean; //hyperprior
  real<lower=0> mu_stdev; //hyperprior
  real sigma_location; //hyperprior
  real<lower=0> sigma_scale; //hyperprior
  //counterfactuals
  int<lower=0> Nx;
  array[Nx] int<lower=1, upper=L> llx;//vec of categories
  array[Nx] row_vector[D] counterfact_x_tilde; // Posterior Prediction intervention
  array[Nx] row_vector[D] counterfact_x; // Posterior Prediction intervention
  //the two statuses to catch relative difference
  array[2] int status_indexes; //the two status indexes to compare (order is x[1] - x[2])
}
parameters {
  array[D] real mu;
  array[D] real<lower=0> sigma;
  array[L] vector[D] beta;
}
model {
  sigma ~ lognormal(sigma_location,sigma_scale); //hyperprior for stdev: shape and inverse scale
  mu ~ normal(mu_mean, mu_stdev); //hyperprior for mean //TODO: convert to mvnormal
  for (l in 1:L) {
    beta[l] ~ normal(mu, sigma);
  }
  {
    vector[N] x_beta_ll;
    for (n in 1:N) {
      x_beta_ll[n] = x[n] * beta[ll[n]];
    }
    y ~ bernoulli_logit(x_beta_ll);
  }
}
generated quantities {
    //SETUP PRIOR PREDICTION
    //preallocate
    real mu_prior[D];
    real sigma_prior[D];
    real p_prior[N]; // what I have priors about
    real p_predicted[N]; //predicted p_values
    //intervention
    real p_predicted_default[Nx]; //predicted p_values
    real p_predicted_intervention[Nx]; //predicted p_values
    real predicted_difference[Nx]; //difference in predicted values

    //collect array of relative status differences between
    array[L] real status_diff;
    //GENERATE RELATIVE DIFFERENCES BETWEEN STATUSES
    for (l in 1:L) {
      status_diff[l] = beta[l,status_indexes[1]] - beta[l,status_indexes[2]];
    }

    //GENERATE PRIOR PREDICTIONS
    {
        vector[D] beta_prior[L];//local var
        //sample parameters
        for (d in 1:D) {
            mu_prior[d] = normal_rng(mu_mean,mu_stdev);
            sigma_prior[d] = lognormal_rng(sigma_location,sigma_scale);
        }
        for (l in 1:L) {
            for (d in 1:D) {
                beta_prior[l,d] = normal_rng(mu_prior[d],sigma_prior[d]);
            }
        }
        //generate probabilities
        vector[D] b_prior[N];//local var
        for (n in 1:N){
            b_prior[n] = beta_prior[ll[n]];
            p_prior[n] = inv_logit(  x[n] * b_prior[n] );
        }
    }
    //GENERATE POSTERIOR PREDICTIONS
    for (n in 1:N) {
        p_predicted[n] = inv_logit( x[n] * beta[ll[n]] );
    }
    //GENERATE POSTERIOR  DISTRIBUTION OF DIFFERENCES
    for (n in 1:Nx) {
        p_predicted_default[n] = inv_logit(  counterfact_x[n] * beta[llx[n]] );
        p_predicted_intervention[n] = inv_logit(  counterfact_x_tilde[n] * beta[llx[n]] ); //intervention

        //intervention - base case
        predicted_difference[n] = p_predicted_intervention[n] - p_predicted_default[n];
    }
}