// Based on lrmc.stan by Ben Goodrich
// This is for the unconstrained partial PO model, which does not handle the
// second set of contrasts for the non-PO part of the model
functions {
// pointwise log-likelihood contributions
vector pw_log_lik(vector alpha, vector beta, matrix tau, vector gamma,
array[] row_vector X, array[] row_vector Z, array[] int y, array[] int cluster) {
int N = size(X);
vector[N] out;
int k = max(y); // assumes all possible categories are observed
int Nc = num_elements(cluster);
int j;
real cj;
real cj1;
for (n in 1:N) {
real eta = X[n] * beta;
if(Nc > 0) eta = eta + gamma[cluster[n]];
j = y[n];
if (j == 1)        cj  = -( alpha[1] + eta );
else if (j == 2)   cj  = alpha[1] + eta;
else               cj  = alpha[j - 1] + eta + Z[n] * tau[ , j - 2];
if(j > 1 && j < k) cj1 = alpha[j] + eta + Z[n] * tau[ , j - 1];
if (j == 1 || j == k) out[n] = log_inv_logit(cj);
else out[n] = log(1./(1. + exp(-cj)) - 1./(1. + exp(-cj1)));
//      else out[n] = log_diff_exp(-log1p_exp(- cj),
//                                 -log1p_exp(- cj1));
//			else out[n] = log(-log1p_exp(-cj) + log1p_exp(-cj1));
}
return out;
}
// Pr(y == j)
matrix Pr(vector alpha, vector beta, matrix tau, vector gamma,
array[] row_vector X, array[] row_vector Z, array[] int y, array[] int cluster) {
int N = size(X);
int k = max(y); // assumes all possible categories are observed
matrix[N, k] out;
for(n in 1:N) {
real eta = X[n] * beta + gamma[cluster[n]];
for(j in 1 : k) {
real cj;
real cj1;
if (j == 1)        cj  = -( alpha[1] + eta );
else if (j == 2)   cj  = alpha[1] + eta;
else               cj  = alpha[j - 1] + eta + Z[n] * tau[ , j - 2];
if(j > 1 && j < k) cj1 = alpha[j] + eta + Z[n] * tau[ , j - 1];
if (j == 1 || j == k) out[n, j] = log_inv_logit(cj);
else out[n, j] = log(1./(1. + exp(-cj)) - 1./(1. + exp(-cj1)));
//				else  out[n, j] = log_diff_exp(-log1p_exp(-cj),
//                                   -log1p_exp(-cj1));
//				else out[n, j] = log(-log1p_exp(-cj) + log1p_exp(-cj1));
}
}
return exp(out);
}
}
data {
int<lower = 1> N;   // number of observations
int<lower = 0> Nc;  // number of clusters  (0 = no clustering)
int<lower = 1> p;   // number of predictors
int<lower = 1> q;   // number of non-PO predictors in Z
int<lower = 0> cn;  // number of contrasts given priors
matrix[N, p] X;     // matrix of CENTERED predictors
matrix[N, q] Z;     // matrix of CENTERED PPO predictors
matrix[cn, p] C;    // contrasts
int<lower = 2> k;   // number of outcome categories
// int<lower = 1, upper = k> y[N]; // outcome on 1 ... k
// int<lower = 1, upper = Nc> cluster[Nc == 0 ? 0 : N];  // cluster IDs
array[N] int<lower = 1, upper = k> y; // outcome on 1 ... k
array[Nc == 0 ? 0 : N] int<lower = 1, upper = Nc> cluster;  // cluster IDs
// prior standard deviations
vector<lower = 0>[q] sdsppo;
// prior means and SDs for contrasts
vector[cn] cmus;
vector[cn] csds;
int<lower = 1, upper = 2> psigma;  // 1=t(4, rsdmean, rsdsd); 2=exponential
// real<lower = 0> rsdmean[Nc == 0 ? 0 : 1];  // mean of prior for sigma
// real<lower = 0> rsdsd[Nc == 0 || psigma == 2 ? 0 : 1];
array[Nc == 0 ? 0 : 1] real<lower = 0> rsdmean;  // mean of prior for sigma
array[Nc == 0 || psigma == 2 ? 0 : 1] real<lower = 0> rsdsd;
// scale parameter for sigma (used only if psigma=1)
real<lower = 0> conc;
}
transformed data {
array[N] row_vector[p] Xr; //row_vector[p] Xr[N];
array[N] row_vector[q] Zr; //row_vector[q] Zr[N];
for (n in 1:N) Xr[n] = X[n, ];
for (n in 1:N) Zr[n] = Z[n, ];
}
parameters {
vector[p] beta; // coefficients on X
matrix[q, k - 2] tau;  // coefficients on Z
simplex[k] pi;  // category probabilities for a person w/ average predictors
vector[Nc] gamma_raw;  // unscaled random effects
// real<lower = 0> sigmag[Nc == 0 ? 0 : 1];   // SD of random effects
array[Nc ==0 ? 0 : 1] real<lower = 0> sigmag;   // SD of random effects
}
transformed parameters {
vector[k - 1] alpha;                               // intercepts
vector[Nc] gamma = Nc == 0 ? gamma_raw : sigmag[1] * gamma_raw;             // scaled random effects
vector[N] log_lik;                                 // log-likelihood pieces
for (j in 2:k) alpha[j - 1] = logit(sum(pi[j:k])); // predictors are CENTERED
log_lik = pw_log_lik(alpha, beta, tau, gamma, Xr, Zr, y, cluster);
}
model {
if(Nc > 0) {
gamma_raw ~ std_normal(); // implies: gamma ~ normal(0, sigmag)
if(psigma == 1) sigmag ~ student_t(4, rsdmean[1], rsdsd[1]);
else sigmag ~ exponential(1. / rsdmean[1]);
}
target += log_lik;
target += dirichlet_lpdf(pi | rep_vector(conc, k));
for (j in 1:(k - 2)) target += normal_lpdf(tau[ , j] | 0, sdsppo);
if(cn > 0) target += normal_lpdf(C * beta | cmus, csds);
}