functions {
  real partial_sum(array[] int slice_obs, int start, int end,
                   array[] int person, array[] int pair, array[] int y,
                   array[] int item_i, array[] int item_k, array[] int trait,
                   matrix tau, vector lambda, vector denom, matrix theta) {

    real lp = 0.0;

    for (obs in start:end) {
      int p = pair[obs];
      int nn = person[obs];
      int i = item_i[p];
      int k = item_k[p];
      int y_cat = y[obs];

      // Calculate the mean of the latent response
      real mu_y = (lambda[i] * theta[nn, trait[i]] - lambda[k] * theta[nn, trait[k]]) / denom[p];

      // Extract the ordered thresholds for this specific pair
      // tau[p] is a row vector containing the C-1 thresholds for pair p
      vector[cols(tau)] thresholds = tau[p]';

      // Scale the thresholds by the denominator to maintain theta parameterization
      vector[cols(tau)] scaled_thresholds = thresholds ./ denom[p];

      // Phi() is the standard normal CDF in Stan.
      // This calculates the log-probability of the Bernoulli outcome safely.
      lp += ordered_probit_lpmf(y_cat | mu_y, scaled_thresholds);
    }

    return lp;
  }
}

data {
  int<lower=1> N; int<lower=1> K; int<lower=1> P; int<lower=1> D; int<lower=1> N_obs;

  int<lower=2> C;          // Number of response categories (e.g., 5)
  int<lower=0, upper=1> has_intransitivity; // 1 if block_size >= 3, else 0

  array[N_obs] int<lower=1, upper=N> person;
  array[N_obs] int<lower=1, upper=P> pair;
  array[N_obs] int<lower=1, upper=C> y;


  array[P] int<lower=1, upper=K> item_i;
  array[P] int<lower=1, upper=K> item_k;
  array[K] int<lower=1, upper=D> trait;
  vector[K] item_sign;

  int<lower=0> K_free;
  array[K] int<lower=0, upper=1> is_fixed_psi;
  vector[K] fixed_psi_val;
  array[K_free] int<lower=1, upper=K> free_psi_idx;
}

transformed data {
  array[N_obs] int obs_idx_array;
  for (i in 1:N_obs) obs_idx_array[i] = i;
  int n_thresh = C - 1; // Number of thresholds is C-1
}

parameters {
  matrix[D, N] z_theta;
  cholesky_factor_corr[D] L_Omega;

  // Matrix of ordered thresholds. P rows (one per pair), n_thresh columns.
  array[P] ordered[n_thresh] tau_raw;

  // Intransitivity error variance (omega^2), constrained to be positive.
  // We declare an array of size 0 or 1. If has_intransitivity == 0, it doesn't exist.
  array[has_intransitivity] real<lower=0> omega2;

  vector<lower=0>[K] lambda_raw;
  vector<lower=0>[K_free] psi2_free;
}

transformed parameters {
  matrix[N, D] theta = (L_Omega * z_theta)';
  vector[K] lambda = lambda_raw .* item_sign;
  vector[K] psi2;
  for (k in 1:K) {
    if (is_fixed_psi[k] == 1) psi2[k] = fixed_psi_val[k];
  }
  for (f in 1:K_free) {
    psi2[free_psi_idx[f]] = psi2_free[f];
  }

  // Calculate the intransitivity variance to add to the denominator
  real omega2_val = 0.0;
  if (has_intransitivity == 1) {
    omega2_val = omega2[1];
  }

  vector[P] denom;
  for (p in 1:P) {
    // Denominator now includes omega2_val if block_size >= 3
    denom[p] = sqrt(psi2[item_i[p]] + psi2[item_k[p]] + omega2_val);
  }

  // Convert the array of ordered thresholds into a matrix
  matrix[P, n_thresh] tau;
  for (p in 1:P) {
    for (c in 1:n_thresh) {
      tau[p, c] = tau_raw[p, c];
    }
  }
}

model {
  // Priors
  to_vector(z_theta) ~ std_normal();
  L_Omega ~ lkj_corr_cholesky(2.0);
  lambda_raw ~ lognormal(0, 0.5);
  psi2_free ~ lognormal(0, 0.5);

  // Prior for the intransitivity variance (if it exists)
  if (has_intransitivity == 1) {
    omega2[1] ~ lognormal(0, 0.5);
  }

  // Threshold priors (normal(0, 3) provides a safe spread)
  for (p in 1:P) {
    tau_raw[p] ~ normal(0, 3);
  }

  // Multithreading Likelihood evaluation
  int grainsize = 250;
  target += reduce_sum(partial_sum, obs_idx_array, grainsize,
                       person, pair, y, item_i, item_k, trait,
                       tau, lambda, denom, theta); // Pass tau instead of alpha
}

generated quantities {
  corr_matrix[D] Omega = multiply_lower_tri_self_transpose(L_Omega);
}


