// gdpar Path 1 -- Hierarchical Bayesian AMM canonical form (v0)
//
// Reference framework, per individual i and (optional) group g[i]:
//   theta_i = theta_ref[g[i]] + a(x_i) + b(x_i) * theta_ref[g[i]]
//             + (W(theta_ref[g[i]]) - W(theta_anchor)) x_i
// (where * denotes elementwise / scalar multiplication for the b term).
//
// p (dimension of theta_ref per group) is fixed at 1 in this version.
// Multi-parameter individual-specific extensions are planned for a
// future version.
//
// 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 scalar theta_ref (matches the Block 6 model up
// to the vector[1] vs real declaration). When use_groups == 1,
// theta_ref is a vector of length J_groups sampled from a hierarchical
// prior Normal(mu_theta_ref, sigma_theta_ref); individuals share
// theta_ref by group. This canonizes the per-group anchor required by
// Block 7 (eBird external validation) without conflating it with the
// additive component a (see (C5) below).
//
// Identifiability condition (C5), Block 6.5: when use_groups == 1, the
// additive design matrix Z_a must not contain a factor(group) column
// or any other rank-deficient indicator aliasing theta_ref[g]. The
// R-side preflight check enforces this; the Stan template assumes the
// spec has passed it.
//
// W basis types supported after sub-phase 8.3.8 (2026-05-22):
//   1 = polynomial  (degrees 1..W_degree of theta_ref, anchored at
//                    theta_anchor)
//   2 = bspline     (Cox-de Boor recursion with fixed boundary_knots
//                    and interior knots; df = W_per_k_dim).
// Dispatch lives in the functions{} block via apply_W_basis_diff(),
// mirroring the apply_inv_link_by_id() canonicalization of sub-phase
// 8.3.7. Generated by R/stan_codegen.R via placeholder substitution;
// do not edit the substituted output by hand.
//
// Family identifiers:
//   1 = gaussian (identity link)
//   2 = poisson (log link)
//   3 = neg_binomial_2 (log link, Stan parametrization with phi)
//   4 = bernoulli (logit link)
//
// Centering of the additive and multiplicative bases is enforced by
// column-wise centering of Z_a and Z_b in the R-side AMM design
// constructor (see amm_spec_designs() / amm_spec_multi_designs()).
// With colMeans(Z_a) = 0, assumption (C2) of Block 1, E_mu[a(X)] = 0,
// holds identically for any a_coef in R^{J_a}; (C3) holds analogously
// for b. No further constraint on a_coef or c_b is required and the
// implementation does not impose one --- doing so would arbitrarily
// restrict the basis coefficients to a (J - 1)-dimensional subspace
// without canonical justification.
//
// Anchoring of the modulating component W via reparametrization
// W(theta) - W(theta_anchor) satisfies assumption (C4) of Block 1
// exactly.
//
// Internal sampling parametrization:
//   a   : configurable CP/NCP. NCP (default): a_raw[j] ~ normal(0, 1)
//         and a_coef[j] = a_raw[j] * sigma_a[1]. CP: a_raw[j] ~
//         normal(0, sigma_a[1]) and a_coef[j] = a_raw[j]. Selected at
//         codegen time via the cp_a argument of generate_stan_code();
//         dispatch in gdpar() is automated by pre-flight diagnostics.
//   W   : configurable CP/NCP. NCP (default): W_raw[k, j] ~
//         normal(0, 1) and the scaling W_raw[k, j] * sigma_W[1] is
//         applied at use site in transformed parameters. CP:
//         W_raw[k, j] ~ normal(0, sigma_W[1]) and no scaling is
//         applied at use site. Selected at codegen time via the
//         cp_W argument of generate_stan_code().
//   b   : LINEAR REPARAMETRIZATION. The multiplicative component
//         enters the linear predictor as theta_ref * (Z_b @ b_coef).
//         Defining c_b = theta_ref * b_coef, the data identify c_b
//         linearly via the term Z_b @ c_b in eta_i. The model
//         samples c_b directly (NCP on c_b_raw), and reports
//         b_coef = c_b / theta_ref as a derived quantity.
// Rationale: when the model is parametrized as (theta_ref, a, b),
// the term theta_ref * (Z_b @ b) makes the log-posterior non-convex
// in the (theta_ref, b) subspace, admitting a spurious local mode
// where theta_ref shifts and b explodes in magnitude with sign
// flipped. This bimodality is intrinsic to the (theta_ref, b)
// parametrization and is independent of NCP vs CP choices.
// Sampling in (theta_ref, a, c_b) yields a strictly log-concave
// (Gaussian-conditional) posterior with a unique mode; the
// pathology is eliminated by construction. Empirical evidence
// (synthetic_recovery Recovery 2, n=1000, b_true=(0.3,-0.3)):
// the (theta_ref, b) parametrization traps ~20-40% of replicates in
// the spurious mode regardless of NCP or CP. The linear
// reparametrization recovers all replicates to the true mode.
// Reference: Stan User's Guide "Reparameterization and Change of
// Variables"; Gelman et al. BDA3 chapter 16 on linear regression
// reparametrization. See refinement of operational decision 7
// (handoff 4 addendum, 2026-05-09).
// AMM canonical form (theta_i = theta_ref + a + b*theta_ref + W) is
// preserved at the user-facing level; the reparametrization is
// internal to the sampler.

functions {
  // {{CANONICAL_HELPERS}}
}

data {
  int<lower=1> n;
  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;

  int<lower=0> J_a;
  int<lower=0> J_b;
  int<lower=0> dim_W;
  int<lower=0> d;

  matrix[n, J_a] Z_a;
  matrix[n, J_b] Z_b;
  matrix[n, d] X;

  vector[n] y_real;
  array[n] int y_int;

  real theta_anchor;

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

  // Sub-phase 8.3.8 (2026-05-22): W basis dispatch field. 0 = W off,
  // 1 = polynomial (anchored at theta_anchor via pow(., jj_local)),
  // 2 = bspline (Cox-de Boor recursion). The auxiliary fields
  // W_n_knots_full, W_knots_full, W_degree are consumed by
  // apply_W_basis_diff() in the functions{} block.
  int<lower=0, upper=2> W_type_id;
  int<lower=0> W_n_knots_full;
  vector[W_n_knots_full] W_knots_full;
  int<lower=0> W_degree;

  // 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].
  int<lower=0, upper=1> use_groups;
  int<lower=1> J_groups;
  array[n] int<lower=1, upper=J_groups> group_id;
}

transformed data {
  int J_a_free = (use_a == 1 && J_a > 0) ? J_a : 0;
  int J_b_free = (use_b == 1 && J_b > 0) ? J_b : 0;
  int dim_W_eff = (use_W == 1) ? dim_W : 0;
  int d_eff = (use_W == 1) ? d : 0;
}

parameters {
  vector[J_groups] 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}}.
  array[use_groups] real mu_theta_ref;
  array[use_groups] real<lower=0> sigma_theta_ref;

  array[use_a == 1 && J_a > 0 ? 1 : 0] real<lower=0> sigma_a;
  vector[J_a_free] a_raw;

  array[use_b == 1 && J_b > 0 ? 1 : 0] real<lower=0> sigma_b;
  vector[J_b_free] c_b_raw;

  array[use_W == 1 && dim_W > 0 ? 1 : 0] real<lower=0> sigma_W;
  matrix[dim_W_eff, d_eff] W_raw;

  array[use_dispersion_y] real<lower=0> sigma_y;
  array[use_dispersion_phi] real<lower=0> phi;
}

transformed parameters {
  vector[J_a] a_coef = rep_vector(0, J_a);
  vector[J_b] c_b = rep_vector(0, J_b);
  vector[J_b] b_coef = rep_vector(0, J_b);
  vector[n] eta;

  if (use_a == 1 && J_a > 0) {
    for (j in 1:J_a_free) {
      a_coef[j] = a_raw[j]{{A_SCALE}};
    }
  }

  if (use_b == 1 && J_b > 0) {
    for (j in 1:J_b_free) {
      c_b[j] = c_b_raw[j] * sigma_b[1];
    }
    // b_coef = c_b / theta_ref 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, so b_coef is
    // left at its initialized zero and users must derive it post-hoc
    // per group from c_b and theta_ref[g] if desired.
    if (use_groups == 0 && theta_ref[1] != 0) {
      for (j in 1:J_b) {
        b_coef[j] = c_b[j] / theta_ref[1];
      }
    }
  }

  for (i in 1:n) {
    real theta_ref_i = theta_ref[group_id[i]];
    real eta_i = theta_ref_i;
    if (use_a == 1 && J_a > 0) {
      eta_i += Z_a[i] * a_coef;
    }
    if (use_b == 1 && J_b > 0) {
      eta_i += Z_b[i] * c_b;
    }
    if (use_W == 1 && dim_W > 0 && d > 0) {
      vector[d] W_diff_x = rep_vector(0, d);
      vector[dim_W] basis_diff = apply_W_basis_diff(
        W_type_id, theta_ref_i, theta_anchor, dim_W, W_degree,
        W_n_knots_full, W_knots_full
      );
      for (k in 1:dim_W) {
        for (jj in 1:d) {
          W_diff_x[jj] += basis_diff[k] * W_raw[k, jj]{{W_SCALE}};
        }
      }
      eta_i += dot_product(W_diff_x, to_vector(X[i]));
    }
    eta[i] = eta_i;
  }
}

model {
  // BEGIN PRIORS
  if (use_groups == 1) {
    mu_theta_ref[1] ~ {{PRIOR_THETA_REF}};
    sigma_theta_ref[1] ~ {{PRIOR_SIGMA_THETA_REF}};
    theta_ref ~ normal(mu_theta_ref[1], sigma_theta_ref[1]);
  } else {
    theta_ref[1] ~ {{PRIOR_THETA_REF}};
  }

  if (use_a == 1 && J_a > 0) {
    sigma_a[1] ~ {{PRIOR_SIGMA_A}};
    a_raw ~ {{A_PRIOR}};
  }
  if (use_b == 1 && J_b > 0) {
    sigma_b[1] ~ {{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[1] ~ {{PRIOR_SIGMA_Y}};
  }
  if (use_dispersion_phi == 1) {
    phi[1] ~ {{PRIOR_PHI}};
  }
  // END PRIORS

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

generated quantities {
  vector[n] log_lik;
  vector[n] theta_i = eta;
  vector[n] y_pred;

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