// extrapolate_epsilon_monotone.stan
// Monotone integrated Gaussian process for the epsilon curve.
// This matches the appendix description:
//   g(x) = softplus(h(x)) >= 0, h ~ GP(0, k)
//   f(x) = f0 + integral g(t) dt  (approximated on a grid)
//   mu_epsilon(x) = L + (epsilon0 - L) * inv_logit(-f(x))
//   y ~ Normal(mu_epsilon, sigma_y)
//
// The construction guarantees a monotone non-increasing epsilon curve because
// f(x) is non-decreasing and inv_logit(-f(x)) is non-increasing.

functions {
  real k_ff(real x, real xp, real sf, real ell) {
    real r2 = (x - xp) * (x - xp);
    return sf * sf * exp(-0.5 * r2 / (ell * ell));
  }
}

data {
  int<lower=1> N;                         // number of observed points
  vector[N] x_obs;                        // observed input locations
  vector[N] y;                            // observed epsilon values

  int<lower=2> M;                         // grid size for integration / prediction
  vector[M] x_grid;                       // strictly increasing grid over the domain
  array[N] int<lower=1, upper=M> obs_idx; // map observed x positions to grid indices

  real<lower=0> range_r;                  // max(x_grid) - min(x_grid)
  real<lower=0> sd_y_prior;               // typical observation-noise scale
  real<lower=0, upper=1> epsilon0;        // left-edge or baseline error rate
}

transformed data {
  vector[M - 1] dx;
  for (m in 1:(M - 1)) dx[m] = x_grid[m + 1] - x_grid[m];
}

parameters {
  real Logit_L;                           // lower asymptote on logit scale
  real<lower=0> sf_h;                     // GP amplitude
  real<lower=1e-3> ell_h;                 // GP length scale
  real<lower=0> sigma_y;                  // observation noise
  real f0;                                // initial latent level
  vector[M] z_h;                          // non-centered latent GP draw
}

transformed parameters {
  real<lower=0, upper=1> L_asym = inv_logit(Logit_L);
  matrix[M, M] K_h;
  matrix[M, M] Lchol_h;
  vector[M] h;
  vector[M] g;
  vector[M] f_grid;
  real jitter;

  jitter = 1e-9 + 1e-8 * sf_h * sf_h;
  for (i in 1:M) {
    for (j in i:M) {
      real v = k_ff(x_grid[i], x_grid[j], sf_h, ell_h);
      if (i == j) v = v + jitter;
      K_h[i, j] = v;
      K_h[j, i] = v;
    }
  }

  Lchol_h = cholesky_decompose(K_h);
  h = Lchol_h * z_h;

  for (m in 1:M) {
    g[m] = log1p_exp(h[m]);
  }

  f_grid[1] = f0;
  for (m in 2:M) {
    f_grid[m] = f_grid[m - 1] + 0.5 * (g[m - 1] + g[m]) * dx[m - 1];
  }
}

model {
  // Hyperpriors chosen to match Appendix E.1 as closely as possible.
  Logit_L ~ normal(-6, 1.3);
  sf_h ~ normal(0, 1.5);
  ell_h ~ lognormal(log(0.15 * range_r), 0.5);
  sigma_y ~ normal(sd_y_prior, sd_y_prior);
  f0 ~ normal(-6, 1.5);
  z_h ~ normal(0, 1);

  for (n in 1:N) {
    real mu_n = L_asym + (epsilon0 - L_asym) * inv_logit(-f_grid[obs_idx[n]]);
    y[n] ~ normal(mu_n, sigma_y);
  }
}

generated quantities {
  vector[M] mu_grid;
  vector[M] y_pred_grid;

  for (m in 1:M) {
    mu_grid[m] = inv_logit(Logit_L) + (epsilon0 - inv_logit(Logit_L)) * inv_logit(-f_grid[m]);
    y_pred_grid[m] = normal_rng(mu_grid[m], sigma_y);
  }
}
