functions {
  real partial_sum(array[] int slice_blocks, int start, int end,
                   array[] int block_person, array[] int block_start_idx, array[] int block_n_pairs,
                   vector y, array[] int pair, array[] int item_i, array[] int item_k, array[] int trait,
                   vector alpha, vector lambda, vector psi2, matrix theta) {

    real lp = 0.0;

    for (b_idx in start:end) {
      int nn = block_person[b_idx];
      int start_idx = block_start_idx[b_idx];
      int n_p = block_n_pairs[b_idx];

      if (n_p == 0) continue;

      vector[n_p] mu_block;
      vector[n_p] y_block;
      matrix[n_p, n_p] Sigma_block;

      Sigma_block = rep_matrix(0.0, n_p, n_p);

      // 1. Build the Mean Vector and the Diagonal of the Covariance Matrix
      for (p in 1:n_p) {
        int obs_idx = start_idx + p - 1; // The row in the flattened vector
        int pair_id = pair[obs_idx];     // The global pair ID (1 to 48)

        int i = item_i[obs_idx];
        int k = item_k[obs_idx];

        // Use pair_id to find the correct intercept in alpha!
        mu_block[p] = alpha[pair_id] + lambda[i] * theta[nn, trait[i]] - lambda[k] * theta[nn, trait[k]];
        y_block[p] = y[obs_idx];

        Sigma_block[p, p] = psi2[i] + psi2[k];
      }

      // 2. Build the Off-Diagonals (Local Dependence)
      if (n_p > 1) {
        for (p1 in 1:(n_p - 1)) {
          for (p2 in (p1 + 1):n_p) {
            int obs1 = start_idx + p1 - 1;
            int obs2 = start_idx + p2 - 1;

            int i1 = item_i[obs1]; int k1 = item_k[obs1];
            int i2 = item_i[obs2]; int k2 = item_k[obs2];

            real cov_val = 0.0;
            if (i1 == i2) cov_val += psi2[i1];
            if (k1 == k2) cov_val += psi2[k1];
            if (i1 == k2) cov_val -= psi2[i1];
            if (k1 == i2) cov_val -= psi2[k1];

            Sigma_block[p1, p2] = cov_val;
            Sigma_block[p2, p1] = cov_val;
          }
        }
      }

      for (p in 1:n_p) Sigma_block[p, p] += 1e-6;

      // 3. Evaluate the Multivariate Normal Likelihood!
      lp += multi_normal_lpdf(y_block | mu_block, Sigma_block);
    }

    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=1> N_blocks;

  array[N_blocks] int<lower=1, upper=N> block_person;
  array[N_blocks] int<lower=1, upper=N_obs> block_start_idx;
  array[N_blocks] int<lower=0, upper=K> block_n_pairs;

  vector[N_obs] y;
  array[N_obs] int<lower=1, upper=P> pair;
  array[N_obs] int<lower=1, upper=K> item_i;
  array[N_obs] 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_blocks] int block_idx_array;
  for (i in 1:N_blocks) block_idx_array[i] = i;
}

parameters {
  matrix[D, N] z_theta;
  cholesky_factor_corr[D] L_Omega;
  vector[P] alpha; // Size is 48

  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];
  }
}

model {
  to_vector(z_theta) ~ std_normal();
  L_Omega ~ lkj_corr_cholesky(2.0);
  alpha ~ normal(0, 3);
  lambda_raw ~ lognormal(0, 0.5);
  psi2_free ~ lognormal(0, 0.5);

  int grainsize = 50;
  target += reduce_sum(partial_sum, block_idx_array, grainsize,
                       block_person, block_start_idx, block_n_pairs,
                       y, pair, item_i, item_k, trait,
                       alpha, lambda, psi2, theta);
}

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