// gdpar Path 1 multivariate -- Hierarchical Bayesian AMM canonical form
// for theta_ref in R^p with coord-wise likelihood factorization.
//
// Reference framework, per individual i, (optional) group g[i], and
// coordinate k = 1..p:
//   theta_i[k] = theta_ref[g[i]][k] + a_k(x_i)
//                + b_k(x_i) * theta_ref[g[i]][k]
//                + (W_k(theta_ref[g[i]]) - W_k(theta_anchor)) * x_i
//
// Grouping (Block 6.5, 2026-05-18+). When use_groups == 0 (default,
// backward-compatible), J_groups == 1 and group_id[i] == 1 for every
// i, reducing to a single global theta_ref of length p (matches Block
// 6 up to the array[1] vs flat declaration). When use_groups == 1,
// theta_ref[g] is sampled coord-wise from Normal(mu_theta_ref[k],
// sigma_theta_ref[k]). Identifiability condition (C5) of Block 6.5:
// Z_a[k] must not contain a factor(group) indicator under
// use_groups == 1. R-side preflight enforces (C5).
//
// Likelihood factorization (homogeneous family across coordinates):
//   p(y_i | theta_i) = prod_{k=1}^p D(y_ik | theta_i[k])
// where D is selected via family_id. This factorization is the
// canonization of Opcion B of the Phase F architectural decision
// (handoff 10 continuation, 2026-05-11): y_i is multivariate in R^p,
// each coordinate independent given theta_i, with cross-dimensional
// coupling carried exclusively by W(theta_ref). Multi-parametric
// families (Opcion A: theta_i parametrizes a single univariate
// distribution via location + log-scale + ...) are deferred to a
// dedicated post-validation block (project memory entry
// project_gdpar_multiparametric_extension_postvalidation).
//
// Generated by R/stan_codegen.R via placeholder substitution; do not
// edit the substituted output by hand.
//
// Family identifiers (homogeneous across all p coordinates):
//   1 = gaussian (identity link)
//   2 = poisson (log link)
//   3 = neg_binomial_2 (log link, Stan's mean-dispersion parametrization)
//   4 = bernoulli (logit link)
//
// Ragged design handling. The additive and multiplicative bases may
// have different column counts per coordinate (J_a_per_k[k],
// J_b_per_k[k]). Stan does not support ragged arrays natively. Two
// flat-pack devices are used:
//   - Padded design matrices: array[p] matrix[n, J_a_max] Z_a; only
//     columns 1..J_a_per_k[k] are referenced for coordinate k. The
//     unused columns are present in storage but never enter the
//     likelihood (the inner loops iterate over 1..J_a_per_k[k]).
//   - Flat-packed coefficients: vector[total_J_a_free] a_raw, with
//     per-coordinate offsets a_raw_offset[k] computed in
//     transformed data. Same convention for c_b_raw.
//
// W is a globally shared parameter. For separable polynomial and
// B-spline bases, the rows of W_raw split into p consecutive blocks
// of W_per_k_dim rows each; block k corresponds to coordinate k of
// theta_ref. Block start row for k is (k-1)*W_per_k_dim + 1.
// Polynomial bases are supported in F.1; B-spline extension is
// deferred (use_W with bspline aborts on the R side in this version).
//
// Internal sampling parametrization (preserved from amm_main.stan):
//   a   : configurable CP/NCP per-coordinate via segment() priors
//         (placeholders DATA_CP_A_PER_K_DECL, TP_A_BLOCK,
//         MODEL_A_BLOCK; resolved by generate_stan_code_multi from
//         the cp_a_per_k logical vector in Phase H.2).
//   W   : configurable CP/NCP per global toggle (placeholders
//         W_SCALE, W_PRIOR).
//   b   : LINEAR reparametrization per coordinate. The multiplicative
//         term enters as theta_ref[k] * (Z_b[k] @ b_coef[k]).
//         Defining c_b[k] = theta_ref[k] * b_coef[k], the data
//         identify c_b[k] linearly via Z_b[k] @ c_b[k] in eta[i, k].
//         The model samples c_b directly (NCP on c_b_raw) and
//         reports b_coef as a derived quantity. This eliminates the
//         bimodality of the (theta_ref, b) parametrization
//         documented in Recovery 2 (handoff 4 addendum, 2026-05-09).
//
// Centering of the additive and multiplicative bases is enforced by
// column-wise centering of Z_a[k] and Z_b[k] in the R-side AMM design
// constructor (see amm_spec_multi_designs()). With colMeans(Z_a[k]) = 0
// for every k, assumption (C2) of Block 1, E_mu[a(X)] = 0, holds
// identically for any a_coef[k] in R^{J_a_per_k[k]}; (C3) holds
// analogously for b. No further constraint on a_coef[k] or c_b[k] is
// required and the implementation does not impose one --- doing so
// would arbitrarily restrict each per-coordinate coefficient vector
// to a (J - 1)-dimensional subspace without canonical justification.
//
// Anchoring of the modulating component W via reparametrization
// W(theta_ref) - W(theta_anchor) satisfies (C4) of Block 1 exactly.
// The cross-dimensional condition C4-bis (Phase H pending) is NOT
// checked by this template; the template assumes the spec has passed
// pre-flight identifiability checks.

data {
  int<lower=1> n;
  int<lower=1> p;
  int<lower=1, upper=4> family_id;

  int<lower=0, upper=1> use_a;
  int<lower=0, upper=1> use_b;
  int<lower=0, upper=1> use_W;

  // Per-coordinate column counts (zero when the component is inactive
  // for that coordinate).
  int<lower=0> J_a_max;
  int<lower=0> J_b_max;
  array[p] int<lower=0> J_a_per_k;
  array[p] int<lower=0> J_b_per_k;

  // Padded design matrices. For coordinate k only columns
  // 1..J_a_per_k[k] (resp. J_b_per_k[k]) are referenced.
  array[p] matrix[n, J_a_max] Z_a;
  array[p] matrix[n, J_b_max] Z_b;

  // Modulating component: W is shared globally across coordinates.
  // W_per_k_dim is dim_W / p for separable polynomial/bspline bases.
  int<lower=0> dim_W;
  int<lower=0> d;
  int<lower=0> W_per_k_dim;
  matrix[n, d] X;

  // Outcomes per coordinate. Both blocks are allocated; only the one
  // matching family_id is consumed.
  matrix[n, p] y_real;
  array[n, p] int y_int;

  vector[p] theta_anchor;

  int<lower=0, upper=1> use_dispersion_y;
  int<lower=0, upper=1> use_dispersion_phi;
  int<lower=0, upper=1> W_is_polynomial;

  // Grouping (Block 6.5). use_groups == 0 (default) collapses to the
  // single-anchor model: J_groups == 1 and group_id[i] == 1 for every
  // observation. use_groups == 1 enables the hierarchical per-group
  // anchor theta_ref[1..J_groups], with coord-wise hyperprior on
  // (mu_theta_ref[k], sigma_theta_ref[k]).
  int<lower=0, upper=1> use_groups;
  int<lower=1> J_groups;
  array[n] int<lower=1, upper=J_groups> group_id;{{DATA_CP_A_PER_K_DECL}}
}

transformed data {
  array[p] int J_a_free;
  array[p] int J_b_free;
  array[p] int a_raw_offset;
  array[p] int c_b_raw_offset;
  int total_J_a_free = 0;
  int total_J_b_free = 0;
  int dim_W_eff = (use_W == 1) ? dim_W : 0;
  int d_eff = (use_W == 1) ? d : 0;
  int sigma_a_size = (use_a == 1) ? p : 0;
  int sigma_b_size = (use_b == 1) ? p : 0;
  int sigma_W_size = (use_W == 1 && dim_W > 0) ? 1 : 0;
  int sigma_y_size = (use_dispersion_y == 1) ? p : 0;
  int phi_size = (use_dispersion_phi == 1) ? p : 0;

  for (k in 1:p) {
    J_a_free[k] = (use_a == 1 && J_a_per_k[k] > 0) ? J_a_per_k[k] : 0;
    J_b_free[k] = (use_b == 1 && J_b_per_k[k] > 0) ? J_b_per_k[k] : 0;
    a_raw_offset[k] = total_J_a_free;
    c_b_raw_offset[k] = total_J_b_free;
    total_J_a_free += J_a_free[k];
    total_J_b_free += J_b_free[k];
  }
}

parameters {
  array[J_groups] vector[p] theta_ref;
  // Hyperparameters of the group-level prior on theta_ref. Allocated
  // only when use_groups == 1; when use_groups == 0 the model samples
  // theta_ref[1] directly from {{PRIOR_THETA_REF}} (coord-wise i.i.d.).
  array[use_groups] vector[p] mu_theta_ref;
  array[use_groups] vector<lower=0>[p] sigma_theta_ref;

  array[sigma_a_size] real<lower=0> sigma_a;
  vector[total_J_a_free] a_raw;

  array[sigma_b_size] real<lower=0> sigma_b;
  vector[total_J_b_free] c_b_raw;

  array[sigma_W_size] real<lower=0> sigma_W;
  matrix[dim_W_eff, d_eff] W_raw;

  array[sigma_y_size] real<lower=0> sigma_y;
  array[phi_size] real<lower=0> phi;
}

transformed parameters {
  array[p] vector[J_a_max] a_coef;
  array[p] vector[J_b_max] c_b;
  array[p] vector[J_b_max] b_coef;
  matrix[n, p] eta;

  for (k in 1:p) {
    a_coef[k] = rep_vector(0, J_a_max);
    c_b[k] = rep_vector(0, J_b_max);
    b_coef[k] = rep_vector(0, J_b_max);
  }

{{TP_A_BLOCK}}

  if (use_b == 1) {
    for (k in 1:p) {
      if (J_b_per_k[k] > 0) {
        for (j in 1:J_b_free[k]) {
          c_b[k][j] = c_b_raw[c_b_raw_offset[k] + j] * sigma_b[k];
        }
        // b_coef[k] = c_b[k] / theta_ref[1][k] is reported as a
        // derived quantity only in the single-anchor regime
        // (use_groups == 0). With per-group anchors there is no
        // canonical scalar to divide by; b_coef stays initialized to
        // zero and users must derive it post-hoc per group from c_b
        // and theta_ref[g][k] if desired.
        if (use_groups == 0 && theta_ref[1][k] != 0) {
          for (j in 1:J_b_per_k[k]) {
            b_coef[k][j] = c_b[k][j] / theta_ref[1][k];
          }
        }
      }
    }
  }

  for (i in 1:n) {
    int g_i = group_id[i];
    for (k in 1:p) {
      real theta_ref_ik = theta_ref[g_i][k];
      real eta_ik = theta_ref_ik;
      if (use_a == 1 && J_a_per_k[k] > 0) {
        eta_ik += dot_product(
          to_vector(Z_a[k][i, 1:J_a_per_k[k]]),
          a_coef[k][1:J_a_per_k[k]]
        );
      }
      if (use_b == 1 && J_b_per_k[k] > 0) {
        eta_ik += dot_product(
          to_vector(Z_b[k][i, 1:J_b_per_k[k]]),
          c_b[k][1:J_b_per_k[k]]
        );
      }
      if (use_W == 1 && dim_W > 0 && d > 0 && W_per_k_dim > 0) {
        vector[d] W_diff_x_k = rep_vector(0, d);
        if (W_is_polynomial == 1) {
          int row_start = (k - 1) * W_per_k_dim + 1;
          for (jj_local in 1:W_per_k_dim) {
            int row = row_start + jj_local - 1;
            real basis_diff = pow(theta_ref_ik, jj_local)
                            - pow(theta_anchor[k], jj_local);
            for (jjj in 1:d) {
              W_diff_x_k[jjj] += basis_diff * W_raw[row, jjj]{{W_SCALE}};
            }
          }
        }
        eta_ik += dot_product(W_diff_x_k, to_vector(X[i]));
      }
      eta[i, k] = eta_ik;
    }
  }
}

model {
  if (use_groups == 1) {
    mu_theta_ref[1] ~ {{PRIOR_THETA_REF}};
    sigma_theta_ref[1] ~ {{PRIOR_SIGMA_THETA_REF}};
    for (g in 1:J_groups) {
      theta_ref[g] ~ normal(mu_theta_ref[1], sigma_theta_ref[1]);
    }
  } else {
    theta_ref[1] ~ {{PRIOR_THETA_REF}};
  }

  if (use_a == 1) {
    sigma_a ~ {{PRIOR_SIGMA_A}};
{{MODEL_A_BLOCK}}
  }
  if (use_b == 1) {
    sigma_b ~ {{PRIOR_SIGMA_B}};
    c_b_raw ~ normal(0, 1);
  }
  if (use_W == 1 && dim_W > 0) {
    sigma_W[1] ~ {{PRIOR_SIGMA_W}};
    to_vector(W_raw) ~ {{W_PRIOR}};
  }
  if (use_dispersion_y == 1) {
    sigma_y ~ {{PRIOR_SIGMA_Y}};
  }
  if (use_dispersion_phi == 1) {
    phi ~ {{PRIOR_PHI}};
  }

  for (k in 1:p) {
    if (family_id == 1) {
      y_real[:, k] ~ normal(eta[:, k], sigma_y[k]);
    } else if (family_id == 2) {
      y_int[:, k] ~ poisson_log(eta[:, k]);
    } else if (family_id == 3) {
      y_int[:, k] ~ neg_binomial_2_log(eta[:, k], phi[k]);
    } else if (family_id == 4) {
      y_int[:, k] ~ bernoulli_logit(eta[:, k]);
    }
  }
}

generated quantities {
  matrix[n, p] log_lik;
  matrix[n, p] theta_i = eta;
  matrix[n, p] y_pred;

  for (i in 1:n) {
    for (k in 1:p) {
      if (family_id == 1) {
        log_lik[i, k] = normal_lpdf(y_real[i, k] | eta[i, k], sigma_y[k]);
        y_pred[i, k] = normal_rng(eta[i, k], sigma_y[k]);
      } else if (family_id == 2) {
        log_lik[i, k] = poisson_log_lpmf(y_int[i, k] | eta[i, k]);
        y_pred[i, k] = poisson_log_rng(eta[i, k]);
      } else if (family_id == 3) {
        log_lik[i, k] = neg_binomial_2_log_lpmf(y_int[i, k] | eta[i, k], phi[k]);
        y_pred[i, k] = neg_binomial_2_log_rng(eta[i, k], phi[k]);
      } else if (family_id == 4) {
        log_lik[i, k] = bernoulli_logit_lpmf(y_int[i, k] | eta[i, k]);
        y_pred[i, k] = bernoulli_logit_rng(eta[i, k]);
      }
    }
  }
}
