// Hand-written PiVo conditioning network + FARGAN-style vocoder step.
// Mirrors onnx/conditioning.onnx and onnx/vocoder_step.onnx operator for operator,
// with weights read straight from those ONNX files. Matrix-vector products run on a
// pluggable kernel: plain JS, or a 472-byte WebAssembly SIMD module.

import { instantiateKernel, wasmSimdSupported } from './accel.ts';
import type { KernelExports } from './accel.ts';
import type { Graph, Tensor } from './onnx-lite.ts';

export type VocoderKind = 'wasm' | 'js';
interface WeightRef {
  off: number;
  nin: number;
  nout: number;
}
type WeightInput = { dims: number[]; data: Tensor['data'] };

// Weight matrices by the MatMul node that consumes them (stable across re-exports).
const VOCODER_MATMULS = {
  frameConv: '/frame_conv/MatMul_output_0',
  frameGlu: '/frame_glu/gate/MatMul_output_0',
  gru1x: '/gru1/MatMul_output_0',
  gru1h: '/gru1/MatMul_1_output_0',
  glu1: '/glu1/gate/MatMul_output_0',
  gru2x: '/gru2/MatMul_output_0',
  gru2h: '/gru2/MatMul_1_output_0',
  glu2: '/glu2/gate/MatMul_output_0',
  gru3x: '/gru3/MatMul_output_0',
  gru3h: '/gru3/MatMul_1_output_0',
  glu3: '/glu3/gate/MatMul_output_0',
  skip: '/skip_dense/MatMul_output_0',
  skipGlu: '/skip_glu/gate/MatMul_output_0',
  out: '/signal_out/MatMul_output_0',
};
const CONDITIONING_MATMULS = { dense1: '/dense1/MatMul_output_0', dense2: '/dense2/MatMul_output_0' };

function matmulWeights<K extends string>(graph: Graph, byOutput: Record<K, string>): Record<K, Tensor> {
  const result = {} as Record<K, Tensor>;
  for (const [key, output] of Object.entries(byOutput) as [K, string][]) {
    const node = graph.nodes.find((n) => n.op === 'MatMul' && n.outputs[0] === output);
    const w = node && graph.initializers.get(node.inputs[1]);
    if (!w || w.dims.length !== 2) throw new Error(`vocoder weight for ${output} not found`);
    result[key] = w;
  }
  return result;
}

// A float32 heap plus a matvec over [nin, nout] row-major weights stored in that heap.
class Heap {
  kind: VocoderKind;
  top: number;
  F: Float32Array;
  memory?: WebAssembly.Memory;
  wasmMatvec?: KernelExports['matvec'];
  constructor(floats: number, kind: VocoderKind) {
    this.kind = kind;
    this.top = 0;
    if (kind === 'wasm') {
      const pages = Math.ceil((floats * 4) / 65536) + 1;
      this.memory = new WebAssembly.Memory({ initial: pages, maximum: pages });
      this.wasmMatvec = instantiateKernel(this.memory).matvec;
      this.F = new Float32Array(this.memory.buffer);
    } else {
      this.F = new Float32Array(floats);
    }
  }
  alloc(n: number): number {
    const off = (this.top + 3) & ~3; // 16-byte alignment for v128 loads
    this.top = off + n;
    if (this.top > this.F.length) throw new Error('vocoder heap exhausted');
    return off;
  }
  view(off: number, n: number): Float32Array {
    return this.F.subarray(off, off + n);
  }
  // Store a weight; JS kernel keeps it transposed ([nout, nin]) for contiguous dot products.
  weight(t: WeightInput): WeightRef {
    const [nin, nout] = t.dims;
    if (nin % 2 || nout % 4) throw new Error(`matvec kernel needs even nin and nout % 4 == 0, got ${nin}x${nout}`);
    const off = this.alloc(nin * nout);
    if (this.kind === 'wasm') {
      this.F.set(t.data, off);
    } else {
      for (let i = 0; i < nin; i++) for (let j = 0; j < nout; j++) this.F[off + j * nin + i] = t.data[i * nout + j];
    }
    return { off, nin, nout };
  }
  matvec(w: WeightRef, x: number, y: number): void {
    if (this.kind === 'wasm') {
      (this.wasmMatvec as KernelExports['matvec'])(w.off * 4, x * 4, y * 4, w.nin, w.nout);
      return;
    }
    const F = this.F;
    const nin = w.nin;
    for (let j = 0; j < w.nout; j++) {
      let s = 0;
      let p = w.off + j * nin;
      for (let i = 0; i < nin; i++) s += F[x + i] * F[p++];
      F[y + j] = s;
    }
  }
}

const sigmoid = (v: number): number => 1 / (1 + Math.exp(-v));
const clip1 = (v: number): number => (v < -1 ? -1 : v > 1 ? 1 : v);

const allocState = (H: Heap) => ({ memory: H.alloc(256), gru1: H.alloc(160), gru2: H.alloc(128), gru3: H.alloc(128), frameState: H.alloc(164) });

function allocScratch(H: Heap) {
  const a = (n: number) => H.alloc(n);
  return {
    x328: a(328), t192: a(192), g192: a(192), c3: a(192),
    in1: a(272), xs1: a(480), hs1: a(480), g160: a(160), c5: a(160),
    in2: a(240), xs2: a(384), hs2: a(384), g128a: a(128), c7: a(128),
    in3: a(208), xs3: a(384), hs3: a(384), g128b: a(128), c9: a(128),
    skipIn: a(688), t128: a(128), g128c: a(128), u128: a(128), o40: a(40),
    condIn: a(32 * 3), dense1: a(64 * 3), conv: a(128), cond: a(320),
  };
}

export class HandVocoder {
  kind: VocoderKind;
  H: Heap;
  W: Record<keyof typeof VOCODER_MATMULS, WeightRef>;
  CW: Record<keyof typeof CONDITIONING_MATMULS, WeightRef>;
  condGainW: Float32Array;
  condGainB: number;
  pitchGainW: Float32Array;
  pitchGainB: Float32Array;
  pitchEmbedding: Float32Array;
  convMat: WeightRef;
  s: ReturnType<typeof allocState>;
  b: ReturnType<typeof allocScratch>;
  pred: Float32Array;
  prev: Float32Array;

  static wasmAvailable(): boolean {
    return wasmSimdSupported();
  }

  constructor(vocoderGraph: Graph, conditioningGraph: Graph, kind: VocoderKind = HandVocoder.wasmAvailable() ? 'wasm' : 'js') {
    const vw = matmulWeights(vocoderGraph, VOCODER_MATMULS);
    const cw = matmulWeights(conditioningGraph, CONDITIONING_MATMULS);
    const init = (g: Graph, name: string): Tensor['data'] => {
      const t = g.initializers.get(name);
      if (!t) throw new Error(`missing initializer ${name}`);
      return t.data;
    };
    let floats = (1 << 14) + 192 * 128;
    for (const t of [...Object.values<Tensor>(vw), ...Object.values<Tensor>(cw)]) floats += t.data.length + 4;
    const H = new Heap(floats, kind);
    this.kind = kind;
    this.H = H;
    this.W = Object.fromEntries(Object.entries(vw).map(([k, t]) => [k, H.weight(t as Tensor)])) as HandVocoder['W'];
    this.CW = Object.fromEntries(Object.entries(cw).map(([k, t]) => [k, H.weight(t as Tensor)])) as HandVocoder['CW'];

    this.condGainW = Float32Array.from(init(vocoderGraph, 'cond_gain.weight'));
    this.condGainB = init(vocoderGraph, 'cond_gain.bias')[0];
    this.pitchGainW = Float32Array.from(init(vocoderGraph, 'pitch_gain.weight')); // [4, 192]
    this.pitchGainB = Float32Array.from(init(vocoderGraph, 'pitch_gain.bias'));
    this.pitchEmbedding = Float32Array.from(init(conditioningGraph, 'pitch_embedding.weight')); // [224, 12]
    const convW = init(conditioningGraph, 'conv.weight'); // [128, 64, 3]
    const convMat = new Float32Array(192 * 128);
    for (let o = 0; o < 128; o++) for (let c = 0; c < 64; c++) for (let f = 0; f < 3; f++) convMat[(f * 64 + c) * 128 + o] = convW[(o * 64 + c) * 3 + f];
    this.convMat = H.weight({ dims: [192, 128], data: convMat });

    // Recurrent state, then scratch (same heap layout as before).
    this.s = allocState(H);
    this.b = allocScratch(H);
    this.pred = new Float32Array(44);
    this.prev = new Float32Array(40);
    this.reset();
  }

  reset(): void {
    const F = this.H.F;
    for (const [name, off] of Object.entries(this.s)) {
      const n = ({ memory: 256, gru1: 160, gru2: 128, gru3: 128, frameState: 164 } as Record<string, number>)[name];
      F.fill(0, off, off + n);
    }
  }

  // context: Float32Array(60) = [previous, current, following] x 20 features; periods: 3 ints.
  // Returns a Float32Array(320) view (4 subframes x 80) valid until the next call.
  conditioning(context: Float32Array, periods: Int32Array): Float32Array {
    const F = this.H.F;
    const b = this.b;
    for (let f = 0; f < 3; f++) {
      const o = b.condIn + f * 32;
      for (let k = 0; k < 20; k++) F[o + k] = context[f * 20 + k];
      const e = (periods[f] - 32) * 12;
      for (let k = 0; k < 12; k++) F[o + 20 + k] = this.pitchEmbedding[e + k];
      this.H.matvec(this.CW.dense1, o, b.dense1 + f * 64);
    }
    for (let i = 0; i < 192; i++) F[b.dense1 + i] = Math.tanh(F[b.dense1 + i]);
    // Conv1d(64 -> 128, kernel 3, no padding) over the 3 frames = one matvec over the
    // concatenated [3 x 64] tanh(dense1) rows with W reordered to [(frame, channel), out].
    this.H.matvec(this.convMat, b.dense1, b.conv);
    for (let o = 0; o < 128; o++) F[b.conv + o] = Math.tanh(F[b.conv + o]);
    this.H.matvec(this.CW.dense2, b.conv, b.cond);
    for (let i = 0; i < 320; i++) F[b.cond + i] = Math.tanh(F[b.cond + i]);
    return this.H.view(b.cond, 320);
  }

  #glu(x: number, gate: number, w: WeightRef, n: number): void {
    const F = this.H.F;
    this.H.matvec(w, x, gate);
    for (let i = 0; i < n; i++) F[x + i] = clip1(F[x + i] * sigmoid(F[gate + i]));
  }

  #gru(input: number, xs: number, hs: number, h: number, wx: WeightRef, wh: WeightRef, n: number): void {
    const F = this.H.F;
    this.H.matvec(wx, input, xs);
    this.H.matvec(wh, h, hs);
    for (let i = 0; i < n; i++) {
      const r = sigmoid(F[hs + i] + F[xs + i]);
      const z = sigmoid(F[hs + n + i] + F[xs + n + i]);
      const c = Math.tanh(F[xs + 2 * n + i] + F[hs + 2 * n + i] * r);
      F[h + i] = (F[h + i] - c) * z + c;
    }
  }

  // One 40-sample subframe. cond: 80 floats (any array), period: int. Writes 40 samples to out[offset..].
  step(cond: Float32Array, condOffset: number, period: number, out: Float32Array, outOffset: number): void {
    const F = this.H.F;
    const s = this.s;
    const b = this.b;
    const W = this.W;
    const H = this.H;

    let gl = this.condGainB;
    for (let i = 0; i < 80; i++) gl += cond[condOffset + i] * this.condGainW[i];
    const gain = Math.fround(Math.exp(Math.fround(gl)));
    const g = Math.fround(gain + 1e-5);

    const pred = this.pred;
    const prev = this.prev;
    for (let k = 0; k < 44; k++) {
      let idx = 256 - period + k - 2;
      if (idx >= 256) idx -= period;
      if (idx < 0) idx = 0;
      pred[k] = clip1(Math.fround(F[s.memory + idx] / g));
    }
    for (let k = 0; k < 40; k++) prev[k] = clip1(Math.fround(F[s.memory + 216 + k] / g));

    // frame_conv input = [frame_state(164), frame_state_out(164)]; frame_state_out = [cond, pred, prev].
    F.copyWithin(b.x328, s.frameState, s.frameState + 164);
    const fso = b.x328 + 164;
    for (let i = 0; i < 80; i++) F[fso + i] = cond[condOffset + i];
    F.set(pred, fso + 80);
    F.set(prev, fso + 124);
    H.matvec(W.frameConv, b.x328, b.c3);
    for (let i = 0; i < 192; i++) F[b.c3 + i] = Math.tanh(F[b.c3 + i]);
    this.#glu(b.c3, b.g192, W.frameGlu, 192);
    F.copyWithin(s.frameState, fso, fso + 164);

    // pitch gains: sigmoid(c3 @ W^T + b), W [4, 192]
    const pg = [0, 0, 0, 0];
    for (let q = 0; q < 4; q++) {
      let acc = this.pitchGainB[q];
      for (let i = 0; i < 192; i++) acc += F[b.c3 + i] * this.pitchGainW[q * 192 + i];
      pg[q] = sigmoid(Math.fround(acc));
    }

    const tail = (dst: number, gate: number): void => {
      for (let k = 0; k < 40; k++) F[dst + k] = gate * pred[k + 2];
      F.set(prev, dst + 40);
    };

    F.copyWithin(b.in1, b.c3, b.c3 + 192);
    tail(b.in1 + 192, pg[0]);
    this.#gru(b.in1, b.xs1, b.hs1, s.gru1, W.gru1x, W.gru1h, 160);
    for (let i = 0; i < 160; i++) F[b.c5 + i] = clip1(F[s.gru1 + i]);
    this.#glu(b.c5, b.g160, W.glu1, 160);

    F.copyWithin(b.in2, b.c5, b.c5 + 160);
    tail(b.in2 + 160, pg[1]);
    this.#gru(b.in2, b.xs2, b.hs2, s.gru2, W.gru2x, W.gru2h, 128);
    for (let i = 0; i < 128; i++) F[b.c7 + i] = clip1(F[s.gru2 + i]);
    this.#glu(b.c7, b.g128a, W.glu2, 128);

    F.copyWithin(b.in3, b.c7, b.c7 + 128);
    tail(b.in3 + 128, pg[2]);
    this.#gru(b.in3, b.xs3, b.hs3, s.gru3, W.gru3x, W.gru3h, 128);
    for (let i = 0; i < 128; i++) F[b.c9 + i] = clip1(F[s.gru3 + i]);
    this.#glu(b.c9, b.g128b, W.glu3, 128);

    // skip = [c5(160), c7(128), c9(128), c3(192), pg3*pred[2:42](40), prev(40)]
    const sk = b.skipIn;
    F.copyWithin(sk, b.c5, b.c5 + 160);
    F.copyWithin(sk + 160, b.c7, b.c7 + 128);
    F.copyWithin(sk + 288, b.c9, b.c9 + 128);
    F.copyWithin(sk + 416, b.c3, b.c3 + 192);
    tail(sk + 608, pg[3]);
    H.matvec(W.skip, sk, b.t128);
    for (let i = 0; i < 128; i++) F[b.t128 + i] = clip1(Math.tanh(F[b.t128 + i]));
    H.matvec(W.skipGlu, b.t128, b.g128c);
    for (let i = 0; i < 128; i++) F[b.u128 + i] = F[b.t128 + i] * sigmoid(F[b.g128c + i]);
    H.matvec(W.out, b.u128, b.o40);

    // memory <- memory[40:] ++ excitation
    F.copyWithin(s.memory, s.memory + 40, s.memory + 256);
    for (let k = 0; k < 40; k++) {
      const e = Math.fround(Math.tanh(F[b.o40 + k]) * gain);
      F[s.memory + 216 + k] = e;
      out[outOffset + k] = e;
    }
  }
}