/**
 * Running moments of "which lines were on the page" against "how much the
 * person went on to use the product", and the ridge solve that turns them into
 * a weight per line.
 *
 * Welford rather than sums of squares, and exponentially weighted rather than
 * flat. The bleed is not a learning rate — with a handful of signups a month a
 * learning rate would just be a way of throwing away most of the evidence. It
 * is insurance: the audience and the market drift, and an observation from two
 * years ago should eventually stop outvoting one from last week. Default it
 * slow, and read `bleed` as "how many observations ago do I stop caring", since
 * the effective sample size this converges to is exactly 1 / (1 - bleed).
 *
 * Everything here is pure and dimension-agnostic. The pool it happens to be
 * measuring lives in `@wellformed/shared`, and the fit that feeds it rows lives
 * next door in `fit.ts`.
 */

/** 1 / (1 - 0.998) ≈ 500 observations of effective memory. */
export const DEFAULT_BLEED = 0.998;

export class Moments {
  /** Effective sample size after bleeding. Not a count — it converges. */
  weight = 0;
  meanR = 0;

  private m2R = 0;
  private readonly meanX: Float64Array;
  private readonly m2XX: Float64Array;
  private readonly m2XR: Float64Array;

  constructor(
    readonly dim: number,
    private readonly bleed: number = DEFAULT_BLEED,
  ) {
    this.meanX = new Float64Array(dim);
    this.m2XX = new Float64Array(dim * dim);
    this.m2XR = new Float64Array(dim);
  }

  /**
   * One page, and what came of it.
   *
   * `x` is the indicator vector over the whole pool — 1 for each line that was
   * on the page, 0 for the rest — and `r` is that visitor's durable use.
   *
   * The co-moment update uses the pre-update deviation on both sides scaled by
   * (W-1)/W, which is West's weighted form written so that it is symmetric by
   * construction rather than only in exact arithmetic. That matters: the solve
   * below assumes a symmetric matrix, and a covariance that drifts a few ulps
   * out of symmetry over ten thousand updates is a bug that only shows up as
   * slightly wrong weights.
   */
  push(x: ArrayLike<number>, r: number): void {
    const d = this.dim;
    if (x.length !== d) throw new Error(`Moments: expected ${d} indicators, got ${x.length}.`);

    // Bleed first, so a fresh observation always arrives at full weight 1.
    const decay = this.bleed;
    if (this.weight > 0) {
      this.m2R *= decay;
      for (let i = 0; i < d; i++) this.m2XR[i] = this.m2XR[i]! * decay;
      for (let i = 0; i < d * d; i++) this.m2XX[i] = this.m2XX[i]! * decay;
    }
    this.weight = decay * this.weight + 1;

    const step = 1 / this.weight;
    const scale = 1 - step;

    const dx = new Float64Array(d);
    for (let i = 0; i < d; i++) dx[i] = x[i]! - this.meanX[i]!;
    const dr = r - this.meanR;

    for (let i = 0; i < d; i++) {
      const di = dx[i]!;
      if (di !== 0) {
        for (let j = i; j < d; j++) {
          const add = scale * di * dx[j]!;
          this.m2XX[i * d + j] = this.m2XX[i * d + j]! + add;
          if (j !== i) this.m2XX[j * d + i] = this.m2XX[i * d + j]!;
        }
      }
      this.m2XR[i] = this.m2XR[i]! + scale * di * dr;
    }
    this.m2R += scale * dr * dr;

    for (let i = 0; i < d; i++) this.meanX[i] = this.meanX[i]! + step * dx[i]!;
    this.meanR += step * dr;
  }

  /** Population covariance of the indicators, row-major `dim × dim`. */
  covXX(): Float64Array {
    return this.scaled(this.m2XX);
  }

  /** Covariance of each indicator with the reward. */
  covXR(): Float64Array {
    return this.scaled(this.m2XR);
  }

  /** Share of pages each item appeared on, under the same weighting. */
  rate(): Float64Array {
    return Float64Array.from(this.meanX);
  }

  varR(): number {
    return this.weight > 0 ? this.m2R / this.weight : 0;
  }

  sdR(): number {
    return Math.sqrt(Math.max(0, this.varR()));
  }

  private scaled(source: Float64Array): Float64Array {
    const out = new Float64Array(source.length);
    if (this.weight <= 0) return out;
    const inv = 1 / this.weight;
    for (let i = 0; i < source.length; i++) out[i] = source[i]! * inv;
    return out;
  }
}

/** A whisper on the diagonal, for conditioning only. Not a shrinkage policy. */
const CONDITIONING = 1e-6;

/**
 * Observations before the latent space is allowed one more dimension.
 *
 * The rank is the model's capacity, so it is the thing that has to be rationed
 * when data is scarce — not the size of the pool, which is fixed by how many
 * lines somebody felt like writing. Twenty observations per direction is
 * deliberately cautious.
 */
const OBSERVATIONS_PER_COMPONENT = 20;

export function rankFor(weight: number, dim: number): number {
  return Math.max(0, Math.min(dim, Math.floor(weight / OBSERVATIONS_PER_COMPONENT)));
}

/**
 * A weight per line, solved through a latent space the data chose.
 *
 * Thirty lines is thirty parameters, and against the signups a year like this
 * produces that is hopeless if they are thirty islands — which is what a
 * per-line average makes them, since it can only ever learn about a line from
 * the visits that line was on. What makes them not islands is the covariance
 * matrix: lines that keep appearing together, and whose pages keep doing
 * similarly well, are related *by the record*, and evidence about one is
 * evidence about the other. That relatedness is discovered here rather than
 * asserted anywhere.
 *
 * The mechanism is `k` steps of conjugate gradient on `covXX · w = covXR`,
 * stopped early. Stopping early is not an approximation of the full solve, it
 * is the model: k steps confine the answer to the Krylov space spanned by
 * {s, Ms, M²s, …}, which is exactly partial least squares with k components —
 * a rank-k latent basis built from the covariance matrix and the reward
 * together, oriented at what actually predicts durable use. Directions the data
 * has not spoken to are not estimated badly, they are not estimated at all.
 *
 * The rank grows with the evidence, so the latent space anneals along with the
 * sampler's confidence. At zero observations k is zero and every weight is
 * zero — neutral, which is where the whole thing is asked to start.
 */
export function latentWeights(
  covXX: Float64Array,
  covXR: Float64Array,
  dim: number,
  components: number,
): Float64Array {
  const w = new Float64Array(dim);
  if (components <= 0 || dim === 0) return w;

  let trace = 0;
  for (let i = 0; i < dim; i++) trace += covXX[i * dim + i]!;
  const scale = trace / dim;
  // Nothing has varied. Every line is equally unproven, and saying so with a
  // vector of zeros is what makes the sampler uniform.
  if (!(scale > 0)) return w;

  const lambda = CONDITIONING * scale;
  const residual = Float64Array.from(covXR);
  const direction = Float64Array.from(covXR);
  let residualNorm = dot(residual, residual, dim);
  if (!(residualNorm > 0)) return w;

  const steps = Math.min(components, dim);
  for (let step = 0; step < steps; step++) {
    const applied = new Float64Array(dim);
    for (let i = 0; i < dim; i++) {
      let sum = lambda * direction[i]!;
      const row = i * dim;
      for (let j = 0; j < dim; j++) sum += covXX[row + j]! * direction[j]!;
      applied[i] = sum;
    }

    const curvature = dot(direction, applied, dim);
    // A direction the design never explored carries no information about the
    // reward. Continuing along it would divide by nothing.
    if (!(curvature > 1e-15)) break;

    const stride = residualNorm / curvature;
    for (let i = 0; i < dim; i++) {
      w[i] = w[i]! + stride * direction[i]!;
      residual[i] = residual[i]! - stride * applied[i]!;
    }

    const nextNorm = dot(residual, residual, dim);
    if (!(nextNorm > 1e-18)) break;
    const momentum = nextNorm / residualNorm;
    for (let i = 0; i < dim; i++) direction[i] = residual[i]! + momentum * direction[i]!;
    residualNorm = nextNorm;
  }

  return w;
}

function dot(a: Float64Array, b: Float64Array, dim: number): number {
  let sum = 0;
  for (let i = 0; i < dim; i++) sum += a[i]! * b[i]!;
  return sum;
}