/**
 * Dataset synthesis. Each run: split every type's snippets into a training
 * set and a placement pool, then place events on a silent timeline. The
 * overlap knob controls how events land: a "clean" event keeps at least one
 * snippet-length of separation from everything placed so far, an
 * "overlapper" is aimed inside the span of an already-placed event. The
 * coincidence knob then decides what fraction of those overlappers land on
 * their host's onset rather than at a random offset within its span, within
 * a window of a few milliseconds, which is the hardest case: two onsets a
 * detector cannot separate in time and a mixture a classifier must take
 * apart. Each placed snippet gets a random gain in dB, so level carries no
 * information about identity. Ground truth is exact by construction.
 */
import { mulberry32, shuffle, type Rng } from './rng.ts';
import type { Dataset, GroundTruthEvent, SnippetLibrary, SynthConfig } from './types.ts';
import { FS } from './types.ts';

/** Split snippet indices per type into [train, pool]. Guarantees at least
 *  one snippet on each side when a type has two or more. */
export function splitSnippets(
  lib: SnippetLibrary,
  trainFraction: number,
  rng: Rng,
): { train: number[][]; pool: number[][] } {
  const byType: number[][] = lib.names.map(() => []);
  lib.snippets.forEach((s, i) => byType[s.typeIndex].push(i));
  const train: number[][] = [];
  const pool: number[][] = [];
  for (const idx of byType) {
    const order = shuffle([...idx], rng);
    let nTrain = Math.round(order.length * trainFraction);
    if (order.length >= 2) nTrain = Math.min(Math.max(nTrain, 1), order.length - 1);
    train.push(order.slice(0, nTrain));
    pool.push(order.slice(nTrain));
  }
  return { train, pool };
}

interface Placement { start: number; typeIndex: number; snippetIndex: number; gain: number; }

export function synthesize(lib: SnippetLibrary, cfg: SynthConfig, runIndex: number): Dataset {
  const rng = mulberry32(cfg.seed + runIndex);
  const { train, pool } = splitSnippets(lib, cfg.trainFraction, rng);
  const snipLen = lib.pre + lib.post;
  const coincidence = cfg.coincidence ?? 0;
  const coincWin = Math.round(((cfg.coincidenceMs ?? 0) / 1000) * FS);
  const n = Math.round(cfg.durationS * FS);
  const margin = snipLen; // keep events off the very edges

  // Build the shuffled list of events to place: (type, snippet from the
  // placement pool, gain). Pools are sampled with replacement so
  // eventsPerType does not depend on how many examples were recorded.
  const wanted: { typeIndex: number; snippetIndex: number; gain: number }[] = [];
  for (let k = 0; k < lib.names.length; k++) {
    if (pool[k].length === 0) continue;
    for (let e = 0; e < cfg.eventsPerType; e++) {
      const snippetIndex = pool[k][Math.floor(rng() * pool[k].length)];
      const gainDb = cfg.gainMinDb * rng();
      wanted.push({ typeIndex: k, snippetIndex, gain: Math.pow(10, gainDb / 20) });
    }
  }
  shuffle(wanted, rng);

  const placed: Placement[] = [];
  const maxTries = 200;
  for (const w of wanted) {
    const overlapper = placed.length > 0 && rng() < cfg.overlap;
    let start = -1;
    if (overlapper) {
      const host = placed[Math.floor(rng() * placed.length)];
      // A coincident overlapper starts within +/- coincidenceMs of its
      // host; an ordinary one anywhere inside the host's span. The
      // coincidence draw is skipped entirely when the knob is off, so a
      // seed that predates it still produces the same dataset.
      const coincident = coincidence > 0 && rng() < coincidence;
      const off = coincident
        ? Math.round((2 * rng() - 1) * coincWin)
        : Math.round((rng() - 0.5) * snipLen);
      const s = host.start + off;
      if (s >= 0 && s + snipLen <= n) start = s;
    }
    if (start < 0) {
      // Clean placement: rejection-sample a start at least snipLen away
      // from every placed event (measured start-to-start, which for equal
      // lengths means no sample overlaps).
      for (let t = 0; t < maxTries; t++) {
        const s = margin + Math.floor(rng() * Math.max(1, n - snipLen - 2 * margin));
        let ok = true;
        for (const p of placed) {
          if (Math.abs(p.start - s) < snipLen) { ok = false; break; }
        }
        if (ok) { start = s; break; }
      }
    }
    if (start < 0) continue; // timeline too crowded; drop the event
    placed.push({ start, ...w });
  }

  const x = new Float32Array(n);
  const truth: GroundTruthEvent[] = [];
  for (const p of placed) {
    const snip = lib.snippets[p.snippetIndex].samples;
    for (let i = 0; i < snip.length && p.start + i < n; i++) x[p.start + i] += snip[i] * p.gain;
    truth.push({ t: (p.start + lib.pre) / FS, typeIndex: p.typeIndex, gain: p.gain });
  }
  truth.sort((a, b) => a.t - b.t);

  // Protect against clipping without rescaling typical data: only scale
  // down when the mix actually exceeds full scale.
  let peak = 0;
  for (let i = 0; i < n; i++) { const a = Math.abs(x[i]); if (a > peak) peak = a; }
  if (peak > 0.99) { const g = 0.99 / peak; for (let i = 0; i < n; i++) x[i] *= g; }

  return { x, fs: FS, truth, trainIndices: train };
}