src/lib/utils/stats.ts

/**
 * Statistical helpers for watermark detectors:
 * - one-proportion z-test (Kirchenbauer, k-SemStamp)
 * - normal CDF / survival function (p-values)
 * - regularized incomplete gamma (TextSeal moment-matched Gamma p-value)
 */

/** Standard normal CDF via Abramowitz-Stegun erf approximation (|err| < 1.5e-7). */
export function normalCdf(x: number): number {
  const t = 1 / (1 + 0.2316419 * Math.abs(x));
  const d = 0.3989422804014327 * Math.exp((-x * x) / 2);
  const poly =
    t * (0.319381530 + t * (-0.356563782 + t * (1.781477937 + t * (-1.821255978 + t * 1.330274429))));
  const p = 1 - d * poly;
  return x >= 0 ? p : 1 - p;
}

/** Upper-tail p-value for a z statistic. */
export function zPValue(z: number): number {
  return 1 - normalCdf(z);
}

/** One-proportion z: (observed - gamma*n) / sqrt(n*gamma*(1-gamma)). */
export function proportionZ(observed: number, n: number, gamma: number): number {
  if (n <= 0) return 0;
  return (observed - gamma * n) / Math.sqrt(n * gamma * (1 - gamma));
}

/** ln Gamma(x) (Lanczos approximation). */
export function logGamma(x: number): number {
  const g = 7;
  const c = [
    0.99999999999980993, 676.5203681218851, -1259.1392167224028, 771.32342877765313,
    -176.61502916214059, 12.507343278686905, -0.13857109526572012, 9.9843695780195716e-6,
    1.5056327351493116e-7,
  ];
  if (x < 0.5) {
    return Math.log(Math.PI / Math.sin(Math.PI * x)) - logGamma(1 - x);
  }
  x -= 1;
  let a = c[0];
  const t = x + g + 0.5;
  for (let i = 1; i < g + 2; i++) a += c[i] / (x + i);
  return 0.5 * Math.log(2 * Math.PI) + (x + 0.5) * Math.log(t) - t + Math.log(a);
}

/**
 * Regularized lower incomplete gamma P(k, x).
 * Series for x < k+1, continued fraction otherwise.
 */
export function lowerGammaP(k: number, x: number): number {
  if (x < 0 || k <= 0) return NaN;
  if (x === 0) return 0;
  const lg = logGamma(k);
  if (x < k + 1) {
    // series expansion
    let sum = 1 / k;
    let term = sum;
    for (let n = 1; n < 500; n++) {
      term *= x / (k + n);
      sum += term;
      if (Math.abs(term) < Math.abs(sum) * 1e-15) break;
    }
    return sum * Math.exp(-x + k * Math.log(x) - lg);
  }
  // continued fraction for Q, then P = 1 - Q
  let b = x + 1 - k;
  let c = 1 / 1e-300;
  let d = 1 / b;
  let h = d;
  for (let i = 1; i < 500; i++) {
    const an = -i * (i - k);
    b += 2;
    d = an * d + b;
    if (Math.abs(d) < 1e-300) d = 1e-300;
    c = b + an / c;
    if (Math.abs(c) < 1e-300) c = 1e-300;
    d = 1 / d;
    const del = d * c;
    h *= del;
    if (Math.abs(del - 1) < 1e-15) break;
  }
  const q = Math.exp(-x + k * Math.log(x) - lg) * h;
  return 1 - q;
}

/**
 * Upper-tail p-value of S under Gamma(shape k, scale theta):
 * p = 1 - F_Gamma(S) = Q(k, S/theta).
 */
export function gammaSurvivalP(s: number, shape: number, scale: number): number {
  if (s <= 0) return 1;
  return 1 - lowerGammaP(shape, s / scale);
}

/**
 * TextSeal moment-matched Gamma null for weighted exponential sums.
 * S = sum(w_i * s_i) where s_i ~ Exp-ish with dual-key variance factor
 * thetaR = alpha^2 + (1-alpha)^2.
 * theta_new = thetaR * sum(w^2)/sum(w); k_new = (sum w)^2 / (thetaR * sum w^2).
 */
export function textsealGammaParams(
  weights: number[],
  alpha: number,
): { shape: number; scale: number } {
  const thetaR = alpha * alpha + (1 - alpha) * (1 - alpha);
  let sw = 0;
  let sw2 = 0;
  for (const w of weights) {
    sw += w;
    sw2 += w * w;
  }
  if (sw <= 0 || sw2 <= 0) return { shape: 1, scale: 1 };
  return {
    scale: (thetaR * sw2) / sw,
    shape: (sw * sw) / (thetaR * sw2),
  };
}

/** Mean of numeric array. */
export function mean(xs: number[]): number {
  return xs.length === 0 ? 0 : xs.reduce((a, b) => a + b, 0) / xs.length;
}