/**
 * One initial condition, on every variant's grid.
 *
 * The host's seeded perturbation is one normal deviate per *grid point*
 * (src/mgpu/noise.ts), so two sessions at different lmax seeded from the same
 * integer do not start from the same field — they start from unrelated fields
 * that merely share a random seed. Comparing them would compare two different
 * problems, and every number the comparison produced would be meaningless.
 *
 * So the field is built once, band-limited at the *coarsest* variant's lmax,
 * and evaluated on each variant's own grid:
 *
 *   1. white noise on the coarsest grid
 *   2. analysed there -> coefficients up to lmax_min
 *   3. zero-padded into each variant's coefficient layout
 *   4. synthesized on that variant's grid
 *
 * Steps 3 and 4 are exact: the field is band-limited at lmax_min, and every
 * variant's band contains that, so each one receives the *same function*
 * sampled where it needs it. Running each model's own `init` on it then leaves
 * every session holding the identical spectral state (zero-padded), which is
 * what makes a pointwise comparison at later times mean something.
 *
 * The coarsest variant gets the projected field too, not the raw white noise
 * it was analysed from — otherwise it alone would start somewhere slightly
 * different from the others.
 */
import { lmIndex, nlmCalc } from '../sht/layout.ts';
import { seededNoise } from '../mgpu/noise.ts';
import type { ModelSession } from '../mgpu/session.ts';

/**
 * Re-index coefficients from a band limit into a wider one's layout, zero-
 * filling the degrees the source does not have. Both layouts are SHTNS
 * m-major with mmax = lmax, so nothing but the index mapping changes.
 */
export function prolongCoeffs(
  q: Float32Array,
  lmaxFrom: number,
  lmaxTo: number,
): Float32Array {
  if (lmaxTo === lmaxFrom) return q;
  if (lmaxTo < lmaxFrom) {
    throw new Error(`prolongCoeffs: cannot widen ${lmaxFrom} into a smaller ${lmaxTo}`);
  }
  const out = new Float32Array(2 * nlmCalc(lmaxTo, lmaxTo));
  for (let m = 0; m <= lmaxFrom; m++) {
    for (let l = m; l <= lmaxFrom; l++) {
      const from = 2 * lmIndex(lmaxFrom, l, m);
      const to = 2 * lmIndex(lmaxTo, l, m);
      out[to] = q[from];
      out[to + 1] = q[from + 1];
    }
  }
  return out;
}

/**
 * The same band-limited perturbation, sampled on each session's grid. Order
 * follows `sessions`. Nothing may be in flight on any session's transform
 * plan — the one-off analys/synth here use the plan's own scratch buffers.
 */
export async function sharedNoise(
  sessions: ModelSession[],
  amp: number,
  seed: number,
): Promise<Float32Array[]> {
  let base = sessions[0];
  for (const s of sessions) if (s.cfg.lmax < base.cfg.lmax) base = s;
  const coeffs = await base.sht.analys(seededNoise(base.npts, amp, seed));
  const out: Float32Array[] = [];
  for (const s of sessions) {
    out.push(await s.sht.synth(prolongCoeffs(coeffs, base.cfg.lmax, s.cfg.lmax)));
  }
  return out;
}