// =============================================================================
// 2D BIPOLAR LATENT TRAIT MODEL V6 - Hierarchical L-R Weights
// =============================================================================
//
// V6 CHANGE: Hierarchical L-R weights varying by country, source, and decade
//   - Replaces global simplex[2] lr_weights with additive logit-scale model
//   - logit_weight[n] = global + country_offset[c] + source_offset[k] + decade_offset[d]
//   - All offsets use non-centered parameterization with estimated sigma hyperparameters
//   - Data-poor contexts shrink toward global mean (~55 new scalar parameters)
//
// PRESERVED FROM V5:
//   - Beta-Binomial expert likelihood with K-scaling
//   - V-Party cultural expansion (5 native items + v2pawelf)
//   - Mean-constituent model (individual party estimates for union members)
//   - Segment-based indexing (S segments, R segment-years)
//   - 3-level hierarchical variance (global -> country -> family)
//   - Random walk dynamics within segments
//   - CDU anchor for scale identification (CDU=1375, not CDU/CSU=211)
//   - Non-centered parameterization for efficiency
//   - Zero-inflation model for manifesto data
//   - Country-item intercepts (coding convention differences across countries)
//   - Binomial-logit likelihood for text data (unchanged)
//
// =============================================================================

data {
  // Segment structure
  int<lower=1> S;                      // Number of segments
  int<lower=1> P;                      // Number of countries
  int<lower=1> R;                      // Total unique segment-year combinations
  int<lower=1> T_year;                 // Total number of years
  array[S] int<lower=1> len_theta_ts;  // Number of years per segment

  // Country membership for each segment
  array[S] int<lower=1, upper=P> segment_country;

  // Segment family data
  int<lower=1> F;                                 // Number of party families
  array[S] int<lower=1, upper=F> segment_family;  // Family index for each segment

  // =========================================================================
  // Text data (manifesto + PolDem)
  // =========================================================================
  int<lower=1> N_man;                  // Number of text observations
  int<lower=1> K_man;                  // Number of unique text items
  array[N_man] int<lower=1, upper=K_man> kk_man;   // Text item index
  array[N_man] int<lower=1, upper=S> ss_man;       // Segment index (first constituent for unions)
  array[N_man] int<lower=1, upper=P> pp_man;       // Country index
  array[N_man] int<lower=0> positive;              // Positive mentions
  array[N_man] int<lower=0> sample;                // Total sample size
  array[N_man] int<lower=1, upper=T_year> year_for_man; // Year index

  // Dimension and direction for text data
  array[N_man] int<lower=1, upper=2> dim_idx_man;  // 1=economic, 2=galtan
  array[N_man] int<lower=-1, upper=1> direction_man; // +1=right/TAN, -1=left/GAL

  // V4: Constituent structure for manifesto observations
  int<lower=1> N_const_man_total;                                    // Total entries in const_rr_man
  array[N_man] int<lower=1> n_const_man;                             // Number of constituents per obs
  array[N_man] int<lower=1> const_offset_man;                        // Offset into const_rr_man
  array[N_const_man_total] int<lower=1, upper=R> const_rr_man;       // Constituent rr indices

  // Country-item-year data (used by zero-inflation model)
  int<lower=1> N_ciy;                  // Number of unique country-item-year combinations
  array[N_man] int<lower=1, upper=N_ciy> ciy_idx;

  // =========================================================================
  // Expert dimension-specific data (V5: integer observations + scale size)
  // =========================================================================
  int<lower=1> N_exp_dim;                  // Number of dimension-specific expert observations
  int<lower=1> K_exp_dim;                  // Number of unique expert items
  array[N_exp_dim] int<lower=1, upper=K_exp_dim> kk_exp_dim;
  array[N_exp_dim] int<lower=1, upper=S> ss_exp_dim;
  array[N_exp_dim] int<lower=1, upper=P> pp_exp_dim;
  array[N_exp_dim] int<lower=0> val_dim_int;           // V5: rounded sum = round(mean * K * n_scale)
  array[N_exp_dim] int<lower=1> n_total_exp_dim;       // V5: K * n_scale (total trials)
  array[N_exp_dim] int<lower=1> n_experts_exp_dim;     // V5: K (number of experts)

  // Dimension index for expert data
  array[N_exp_dim] int<lower=1, upper=2> dim_idx_exp;  // 1=economic, 2=galtan

  // V4: Constituent structure for expert dim observations
  int<lower=1> N_const_exp_dim_total;
  array[N_exp_dim] int<lower=1> n_const_exp_dim;
  array[N_exp_dim] int<lower=1> const_offset_exp_dim;
  array[N_const_exp_dim_total] int<lower=1, upper=R> const_rr_exp_dim;

  // =========================================================================
  // Expert general L-R data (V5: integer observations + scale size)
  // =========================================================================
  int<lower=1> N_exp_lr;
  int<lower=1> K_exp_lr;
  array[N_exp_lr] int<lower=1, upper=K_exp_lr> kk_exp_lr;
  array[N_exp_lr] int<lower=1, upper=S> ss_exp_lr;
  array[N_exp_lr] int<lower=1, upper=P> pp_exp_lr;
  array[N_exp_lr] int<lower=0> val_lr_int;             // V5: rounded sum = round(mean * K * n_scale)
  array[N_exp_lr] int<lower=1> n_total_exp_lr;         // V5: K * n_scale (total trials)
  array[N_exp_lr] int<lower=1> n_experts_exp_lr;       // V5: K (number of experts)

  // V6: Decade indexing for hierarchical L-R weights
  int<lower=1> D_lr;                                   // Number of decades
  array[N_exp_lr] int<lower=1, upper=D_lr> dd_exp_lr;  // Decade index per LR obs

  // V4: Constituent structure for expert L-R observations
  int<lower=1> N_const_exp_lr_total;
  array[N_exp_lr] int<lower=1> n_const_exp_lr;
  array[N_exp_lr] int<lower=1> const_offset_exp_lr;
  array[N_const_exp_lr_total] int<lower=1, upper=R> const_rr_exp_lr;

  // Prior information
  real mn_resp_log_man;
  real mn_resp_log_exp_dim;
  real mn_resp_log_exp_lr;

  // Identification anchor
  int<lower=1, upper=S> anchor_segment;  // CDU segment (1375 with unions, 211 without)

  // V3 compatibility: these are still passed but not used in V4+ likelihood
  // (kept so older data dicts work without modification for backwards compat)
  array[N_man] int<lower=1, upper=R> rr_man;       // Segment-year index (unused in V4+ likelihood)
  array[N_exp_dim] int<lower=1, upper=R> rr_exp_dim;
  array[N_exp_lr] int<lower=1, upper=R> rr_exp_lr;
}

transformed data {
  real eps = 1e-6;
  real one_minus_eps = 1 - eps;

  // Count zero and non-zero samples
  int N_man_zero = 0;
  int N_man_nonzero = 0;
  for (n in 1:N_man) {
    if (sample[n] == 0) {
      N_man_zero += 1;
    } else {
      N_man_nonzero += 1;
    }
  }

  // Create index arrays for zero/nonzero split
  array[N_man_zero > 0 ? N_man_zero : 1] int idx_zero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int idx_nonzero;

  {
    int pos_zero = 1;
    int pos_nonzero = 1;
    for (n in 1:N_man) {
      if (sample[n] == 0) {
        idx_zero[pos_zero] = n;
        pos_zero += 1;
      } else {
        idx_nonzero[pos_nonzero] = n;
        pos_nonzero += 1;
      }
    }
  }

  // V4: Pre-compute constituent info for nonzero manifesto obs
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int kk_man_nonzero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int orig_idx_nonzero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int direction_nonzero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int pp_man_nonzero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int dim_idx_nonzero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int n_const_nonzero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int const_offset_nonzero;

  {
    for (i in 1:N_man_nonzero) {
      int n = idx_nonzero[i];
      orig_idx_nonzero[i] = n;
      kk_man_nonzero[i] = kk_man[n];
      direction_nonzero[i] = direction_man[n];
      pp_man_nonzero[i] = pp_man[n];
      dim_idx_nonzero[i] = dim_idx_man[n];
      n_const_nonzero[i] = n_const_man[n];
      const_offset_nonzero[i] = const_offset_man[n];
    }
  }

  // Pre-compute segment start positions
  array[S + 1] int segment_start;
  segment_start[1] = 1;
  for (s in 1:S) {
    segment_start[s + 1] = segment_start[s] + len_theta_ts[s];
  }

  // Extract sample and positive for nonzero observations
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int sample_nonzero;
  array[N_man_nonzero > 0 ? N_man_nonzero : 1] int positive_nonzero;
  for (i in 1:N_man_nonzero) {
    sample_nonzero[i] = sample[idx_nonzero[i]];
    positive_nonzero[i] = positive[idx_nonzero[i]];
  }

  // Pre-compute family-to-country mapping
  array[F] int family_country;
  {
    for (f in 1:F) {
      family_country[f] = 0;
    }
    for (s in 1:S) {
      int f = segment_family[s];
      if (family_country[f] == 0) {
        family_country[f] = segment_country[s];
      }
    }
    for (f in 1:F) {
      if (family_country[f] == 0) {
        family_country[f] = 1;
      }
    }
  }
}

parameters {
  // =========================================================================
  // Latent position parameters - 2 dimensions
  // =========================================================================
  matrix[2, R] theta_ncp;                 // Non-centered: [1]=economic_lr, [2]=galtan
  matrix[2, S] theta_init_raw;            // Initial traits per segment
  vector<lower=0>[2] sigma_theta_init;    // SD of initial theta per dimension

  // =========================================================================
  // Three-level hierarchical variance
  // =========================================================================
  vector[2] mu_sigma_global_raw;          // Global mean RW variance
  vector<lower=0>[2] tau_sigma_country;   // Country deviation scale
  matrix[2, P] sigma_country_raw;         // Non-centered country deviations
  vector<lower=0>[2] tau_sigma_family;    // Family deviation scale
  matrix[2, F] sigma_family_raw;          // Non-centered family deviations

  // =========================================================================
  // Country-item intercepts (coding convention differences)
  // =========================================================================
  matrix[P, K_man] country_item_raw;
  real<lower=0> sigma_country_item;

  // =========================================================================
  // Zero-sample parameters
  // =========================================================================
  real alpha_zs;
  vector[T_year] year_effect_raw;
  real<lower=0> sigma_year_effect;
  vector[S] segment_zs_raw;
  real<lower=0> sigma_segment_zs;
  vector[N_ciy] ciy_zs_raw;
  real<lower=0> sigma_ciy_zs;

  // =========================================================================
  // Item parameters
  // =========================================================================

  // Text data: intercept + single positive loading (direction handled in data)
  vector[K_man] gamma_man_intercept_raw;
  vector<lower=0>[K_man] gamma_man_loading;  // Positive loading

  // Expert dimension-specific: intercept + slope (direct mapping to dimension)
  vector[K_exp_dim] gamma_exp_intercept_raw;
  vector<lower=0>[K_exp_dim] gamma_exp_slope;

  // Expert general L-R: intercept + slope
  vector[K_exp_lr] gamma_lr_intercept_raw;
  vector<lower=0>[K_exp_lr] gamma_lr_slope;

  // V6: Hierarchical L-R weights (replace simplex[2] lr_weights)
  real lr_weight_global;                       // Global logit-scale weight
  vector[P] lr_country_offset_raw;             // Country offsets (non-centered)
  vector[K_exp_lr] lr_source_offset_raw;       // Source offsets (non-centered)
  vector[D_lr] lr_decade_offset_raw;           // Decade offsets (non-centered)
  real<lower=0> sigma_lr_country;              // SD of country offsets
  real<lower=0> sigma_lr_source;               // SD of source offsets
  real<lower=0> sigma_lr_decade;               // SD of decade offsets

  // Precision parameters
  real<lower=0> phi_exp_dim;  // Precision for dimension-specific (now: only measurement noise)
  real<lower=0> phi_exp_lr;   // Precision for general L-R (now: only measurement noise)

  // Scale parameters
  real<lower=0> sigma_intercept_man;
  real<lower=0> sigma_loading_man;
  real<lower=0> sigma_intercept_exp;
  real<lower=0> sigma_slope_exp;
  real<lower=0> sigma_intercept_lr;
  real<lower=0> sigma_slope_lr;
  real mu_lambda_man;
  real mu_lambda_exp_dim;
  real mu_lambda_exp_lr;
}

transformed parameters {
  matrix[2, R] theta;                  // [1]=economic_lr, [2]=galtan
  matrix[2, S] theta_init;

  // Three-level hierarchical variance (2D)
  matrix[2, P] sigma_theta_country;
  for (d in 1:2) {
    for (p in 1:P) {
      sigma_theta_country[d, p] = log1p_exp(
        mu_sigma_global_raw[d] + tau_sigma_country[d] * sigma_country_raw[d, p]
      );
    }
  }

  matrix[2, F] sigma_theta_family;
  for (d in 1:2) {
    for (f in 1:F) {
      int c = family_country[f];
      sigma_theta_family[d, f] = log1p_exp(
        log(sigma_theta_country[d, c]) + tau_sigma_family[d] * sigma_family_raw[d, f]
      );
    }
  }

  // Non-centered parameterization for theta_init
  for (dim in 1:2) {
    theta_init[dim, :] = sigma_theta_init[dim] * theta_init_raw[dim, :];
  }

  // Country-item intercepts
  matrix[P, K_man] country_item_intercept = sigma_country_item * country_item_raw;

  // Zero-sample components
  vector[T_year] year_effect = sigma_year_effect * year_effect_raw;
  vector[S] segment_zs = sigma_segment_zs * segment_zs_raw;
  vector[N_ciy] ciy_zs = sigma_ciy_zs * ciy_zs_raw;

  // Zero-sample probability
  vector[N_man] zero_sample_logit = alpha_zs +
                                     year_effect[year_for_man] +
                                     segment_zs[ss_man] +
                                     ciy_zs[ciy_idx];
  vector[N_man] zero_sample_prob = inv_logit(zero_sample_logit);
  zero_sample_prob = fmax(fmin(zero_sample_prob, one_minus_eps), eps);

  // Construct theta using family-specific random walk variance (2D)
  for (dim in 1:2) {
    for (s in 1:S) {
      int start = segment_start[s];
      int Ts = len_theta_ts[s];
      int fam = segment_family[s];
      real sigma_s = sigma_theta_family[dim, fam];

      theta[dim, start] = theta_init[dim, s] + sigma_s * theta_ncp[dim, start];
      if (Ts > 1) {
        theta[dim, start + 1 : start + Ts - 1] = theta[dim, start] +
          cumulative_sum(sigma_s * theta_ncp[dim, start + 1 : start + Ts - 1]);
      }
    }
  }

  // Item parameters
  vector[K_man] gamma_man_intercept = mu_lambda_man + sigma_intercept_man * gamma_man_intercept_raw;
  vector[K_exp_dim] gamma_exp_intercept = mu_lambda_exp_dim + sigma_intercept_exp * gamma_exp_intercept_raw;
  vector[K_exp_lr] gamma_lr_intercept = mu_lambda_exp_lr + sigma_intercept_lr * gamma_lr_intercept_raw;
}

model {
  // =========================================================================
  // PRIORS
  // =========================================================================

  // Three-level hierarchical variance priors (2D)
  mu_sigma_global_raw ~ normal(-0.8, 0.5);
  tau_sigma_country ~ normal(0, 0.3);
  to_vector(sigma_country_raw) ~ std_normal();
  tau_sigma_family ~ normal(0, 0.2);
  to_vector(sigma_family_raw) ~ std_normal();

  // Other variance priors
  sigma_theta_init ~ normal(0, 0.5);
  to_vector(theta_init_raw) ~ std_normal();
  to_vector(theta_ncp) ~ std_normal();

  // Country-item intercept priors
  to_vector(country_item_raw) ~ std_normal();
  sigma_country_item ~ normal(0, 0.3);

  // Zero-sample priors
  alpha_zs ~ normal(-1, 1);
  year_effect_raw ~ std_normal();
  sigma_year_effect ~ normal(0, 0.5);
  segment_zs_raw ~ std_normal();
  sigma_segment_zs ~ normal(0, 0.3);
  ciy_zs_raw ~ std_normal();
  sigma_ciy_zs ~ normal(0, 0.3);

  // =========================================================================
  // Item parameter priors
  // =========================================================================

  // Text data item priors
  mu_lambda_man ~ normal(mn_resp_log_man, 0.5);
  sigma_intercept_man ~ normal(0, 1);
  sigma_loading_man ~ normal(0, 0.5);
  gamma_man_intercept_raw ~ std_normal();
  gamma_man_loading ~ normal(1.0, sigma_loading_man);

  // Expert dimension item priors
  mu_lambda_exp_dim ~ normal(mn_resp_log_exp_dim, 0.5);
  sigma_intercept_exp ~ normal(0, 1);
  sigma_slope_exp ~ normal(0, 0.5);
  gamma_exp_intercept_raw ~ std_normal();
  gamma_exp_slope ~ normal(1.0, sigma_slope_exp);

  // Expert L-R item priors
  mu_lambda_exp_lr ~ normal(mn_resp_log_exp_lr, 0.5);
  sigma_intercept_lr ~ normal(0, 1);
  sigma_slope_lr ~ normal(0, 0.5);
  gamma_lr_intercept_raw ~ std_normal();
  gamma_lr_slope ~ normal(1.0, sigma_slope_lr);

  // V6: Hierarchical L-R weight priors
  lr_weight_global ~ normal(0, 1);
  lr_country_offset_raw ~ std_normal();
  lr_source_offset_raw ~ std_normal();
  lr_decade_offset_raw ~ std_normal();
  sigma_lr_country ~ normal(0, 0.5);
  sigma_lr_source ~ normal(0, 0.5);
  sigma_lr_decade ~ normal(0, 0.5);

  // Expert data precision priors
  phi_exp_dim ~ gamma(50, 0.5);
  phi_exp_lr ~ gamma(10, 0.5);

  // =========================================================================
  // CDU ANCHOR CONSTRAINT
  // With unions: anchors CDU (1375) at moderate center-right position
  // Without unions: anchors CDU/CSU (211) as before
  // =========================================================================
  target += normal_lpdf(theta_init[1, anchor_segment] | 0.2, 0.2);  // economic_lr
  target += normal_lpdf(theta_init[2, anchor_segment] | 0.2, 0.2);  // galtan

  // =========================================================================
  // LIKELIHOOD 1: Text data (binomial with zero-inflation)
  // V4: Mean-constituent averaging for union observations
  // =========================================================================

  // Zero-sample observations
  if (N_man_zero > 0) {
    target += sum(log(zero_sample_prob[idx_zero]));
  }

  // Non-zero observations with constituent averaging
  if (N_man_nonzero > 0) {
    vector[N_man_nonzero] lin_man;

    for (i in 1:N_man_nonzero) {
      int nc = n_const_nonzero[i];
      int off = const_offset_nonzero[i];
      int dim = dim_idx_nonzero[i];
      real avg_pos;

      if (nc == 1) {
        // Fast path: single party (>95% of observations)
        avg_pos = theta[dim, const_rr_man[off]];
      } else {
        // Union: average over constituent thetas
        avg_pos = 0;
        for (c in 0:(nc-1)) {
          avg_pos += theta[dim, const_rr_man[off + c]];
        }
        avg_pos /= nc;
      }

      lin_man[i] = gamma_man_intercept[kk_man_nonzero[i]] +
                   direction_nonzero[i] * gamma_man_loading[kk_man_nonzero[i]] * avg_pos +
                   country_item_intercept[pp_man_nonzero[i], kk_man_nonzero[i]];
    }

    for (i in 1:N_man_nonzero) {
      target += log1m(zero_sample_prob[orig_idx_nonzero[i]]) +
                binomial_logit_lpmf(positive_nonzero[i] | sample_nonzero[i], lin_man[i]);
    }
  }

  // =========================================================================
  // LIKELIHOOD 2: Expert dimension-specific data (V5: beta-binomial likelihood)
  // V4: Constituent averaging for union-level expert obs
  // =========================================================================
  {
    vector[N_exp_dim] pos;
    for (n in 1:N_exp_dim) {
      int nc = n_const_exp_dim[n];
      int off = const_offset_exp_dim[n];
      int dim = dim_idx_exp[n];

      if (nc == 1) {
        pos[n] = theta[dim, const_rr_exp_dim[off]];
      } else {
        pos[n] = 0;
        for (c in 0:(nc-1)) {
          pos[n] += theta[dim, const_rr_exp_dim[off + c]];
        }
        pos[n] /= nc;
      }
    }

    vector[N_exp_dim] lin_exp_dim;
    for (n in 1:N_exp_dim) {
      lin_exp_dim[n] = gamma_exp_intercept[kk_exp_dim[n]] +
                       gamma_exp_slope[kk_exp_dim[n]] * pos[n];
    }

    vector[N_exp_dim] mu_exp_dim = inv_logit(lin_exp_dim);
    mu_exp_dim = fmax(fmin(mu_exp_dim, one_minus_eps), eps);

    // K-scaling: phi * K corrects for independent expert perceptions
    vector[N_exp_dim] alpha_exp_dim = phi_exp_dim * to_vector(n_experts_exp_dim) .* mu_exp_dim;
    vector[N_exp_dim] beta_exp_dim = phi_exp_dim * to_vector(n_experts_exp_dim) .* (1 - mu_exp_dim);
    val_dim_int ~ beta_binomial(n_total_exp_dim, alpha_exp_dim, beta_exp_dim);
  }

  // =========================================================================
  // LIKELIHOOD 3: Expert general L-R data (V6: hierarchical per-obs weights)
  // V4: Constituent averaging + weighted combination of both dimensions
  // V6: Per-observation weights via country + source + decade offsets
  // =========================================================================
  {
    // V6: Compute per-observation economic weight on logit scale
    vector[N_exp_lr] logit_w;
    for (n in 1:N_exp_lr) {
      logit_w[n] = lr_weight_global
                 + sigma_lr_country * lr_country_offset_raw[pp_exp_lr[n]]
                 + sigma_lr_source  * lr_source_offset_raw[kk_exp_lr[n]]
                 + sigma_lr_decade  * lr_decade_offset_raw[dd_exp_lr[n]];
    }
    vector[N_exp_lr] w_econ = inv_logit(logit_w);

    vector[N_exp_lr] combined_pos;
    for (n in 1:N_exp_lr) {
      int nc = n_const_exp_lr[n];
      int off = const_offset_exp_lr[n];

      if (nc == 1) {
        int r = const_rr_exp_lr[off];
        combined_pos[n] = w_econ[n] * theta[1, r] + (1 - w_econ[n]) * theta[2, r];
      } else {
        // Average the combined position across constituents
        combined_pos[n] = 0;
        for (c in 0:(nc-1)) {
          int r = const_rr_exp_lr[off + c];
          combined_pos[n] += w_econ[n] * theta[1, r] + (1 - w_econ[n]) * theta[2, r];
        }
        combined_pos[n] /= nc;
      }
    }

    vector[N_exp_lr] lin_exp_lr;
    for (n in 1:N_exp_lr) {
      lin_exp_lr[n] = gamma_lr_intercept[kk_exp_lr[n]] +
                      gamma_lr_slope[kk_exp_lr[n]] * combined_pos[n];
    }

    vector[N_exp_lr] mu_exp_lr = inv_logit(lin_exp_lr);
    mu_exp_lr = fmax(fmin(mu_exp_lr, one_minus_eps), eps);

    // K-scaling: phi * K corrects for independent expert perceptions
    vector[N_exp_lr] alpha_exp_lr = phi_exp_lr * to_vector(n_experts_exp_lr) .* mu_exp_lr;
    vector[N_exp_lr] beta_exp_lr = phi_exp_lr * to_vector(n_experts_exp_lr) .* (1 - mu_exp_lr);
    val_lr_int ~ beta_binomial(n_total_exp_lr, alpha_exp_lr, beta_exp_lr);
  }
}

generated quantities {
  // =========================================================================
  // Direct outputs (no derivation needed)
  // =========================================================================

  // Two bipolar scales on [0,1] via inv_logit
  vector[R] economic_lr = inv_logit(to_vector(theta[1, :]));  // 0=left, 1=right
  vector[R] galtan = inv_logit(to_vector(theta[2, :]));       // 0=GAL, 1=TAN

  // General left-right combining both dimensions (using global weight only)
  // Per-R country/source/decade info not available in GQ, so use global mean
  real lr_w_econ_global = inv_logit(lr_weight_global);
  vector[R] general_lr = inv_logit(
    lr_w_econ_global * to_vector(theta[1, :]) + (1 - lr_w_econ_global) * to_vector(theta[2, :])
  );

  // V6: Expose hierarchical L-R weight diagnostics
  real lr_weight_econ_global = lr_w_econ_global;
  real lr_sigma_country = sigma_lr_country;
  real lr_sigma_source = sigma_lr_source;
  real lr_sigma_decade = sigma_lr_decade;

  // Expose hierarchical variance diagnostics (2D)
  vector[2] mean_sigma_global;
  vector[2] mean_sigma_country;
  vector[2] mean_sigma_family;
  for (d in 1:2) {
    mean_sigma_global[d] = log1p_exp(mu_sigma_global_raw[d]);
    mean_sigma_country[d] = mean(sigma_theta_country[d, :]);
    mean_sigma_family[d] = mean(sigma_theta_family[d, :]);
  }
}