/** Dense FP32 inference for the small Motion MLP. Feature construction is a separate layer. */
export class DenseNetwork {
  constructor(manifest, buffer) {
    if (!['superhq-dense-f32-v1', 'superhq-dense-f32-v2', 'superhq-dense-f32-v3'].includes(manifest.format) || manifest.byte_order !== 'little' ||
        !(buffer instanceof ArrayBuffer) || buffer.byteLength !== manifest.bytes || buffer.byteLength % 4) {
      throw new Error('Invalid dense network manifest or weight buffer.');
    }
    const bytes = new Uint8Array(new Uint32Array([1]).buffer);
    if (bytes[0] !== 1) throw new Error('This runtime requires a little-endian platform.');
    this.manifest = manifest;
    const weights = new Float32Array(buffer);
    if (!weights.every(Number.isFinite)) throw new Error('Nonfinite network weights.');
    let offset = 0;
    this.layers = manifest.layers.map((layer, i) => {
      const { inputs, outputs, weight_offset, bias_offset, activation } = layer;
      if (!Number.isInteger(inputs) || !Number.isInteger(outputs) || inputs < 1 || outputs < 1 ||
          weight_offset !== offset || bias_offset !== offset + inputs * outputs ||
          (i > 0 && inputs !== manifest.layers[i - 1].outputs) ||
          !['silu', 'identity'].includes(activation)) throw new Error('Invalid dense layer layout.');
      offset = bias_offset + outputs;
      if (offset > weights.length) throw new Error('Layer exceeds weight buffer.');
      return { inputs, outputs, activation, weight: weights.subarray(weight_offset, bias_offset),
        bias: weights.subarray(bias_offset, offset), output: new Float32Array(outputs) };
    });
    if (!this.layers.length || offset !== weights.length || offset !== manifest.parameters) {
      throw new Error('Weight count differs from manifest.');
    }
    this.outputSkip = manifest.output_skip ?? [];
    this.angularPrior = manifest.angular_prior;
    if ((manifest.format === 'superhq-dense-f32-v3') !== (this.angularPrior === 'articulated_constant_velocity_v1') ||
        (this.angularPrior && (this.layers[0].inputs !== 850 || this.layers.at(-1).outputs !== 144))) {
      throw new Error('Invalid angular prior contract.');
    }
    if (!Array.isArray(this.outputSkip) ||
        (manifest.format === 'superhq-dense-f32-v2' && !this.outputSkip.length) ||
        (manifest.format === 'superhq-dense-f32-v1' && this.outputSkip.length)) throw new Error('Invalid output skip layout.');
    for (const item of this.outputSkip) {
      if (!Number.isInteger(item.input) || item.input < 0 || item.input >= this.layers[0].inputs ||
          !Number.isInteger(item.output) || item.output < 0 || item.output >= this.layers.at(-1).outputs ||
          !Number.isFinite(item.scale)) throw new Error('Invalid output skip connection.');
    }
  }

  /** The returned scratch buffer is reused by the next call. Copy it when retaining a prediction. */
  forward(input) {
    if (!(input instanceof Float32Array) || input.length !== this.layers[0].inputs || !input.every(Number.isFinite)) {
      throw new Error('Expected a finite Float32Array of the network input width.');
    }
    let value = input;
    for (const layer of this.layers) {
      const { inputs, outputs, weight, bias, output, activation } = layer;
      for (let row = 0; row < outputs; row++) {
        let sum = bias[row];
        const start = row * inputs;
        for (let col = 0; col < inputs; col++) sum += weight[start + col] * value[col];
        output[row] = activation === 'silu' ? sum / (1 + Math.exp(-sum)) : sum;
      }
      if (!output.every(Number.isFinite)) throw new Error('Nonfinite network output.');
      value = output;
    }
    for (const item of this.outputSkip) value[item.output] += input[item.input] * item.scale;
    if (this.angularPrior) addAngularPrior(input, value);
    if (!value.every(Number.isFinite)) throw new Error('Nonfinite network output after skip connections.');
    return value;
  }
}

function rotation(input, offset) {
  const a = Array.from(input.subarray(offset, offset + 3));
  const b = Array.from(input.subarray(offset + 3, offset + 6));
  const na = Math.hypot(...a);
  if (na < 1e-6) throw new Error('Degenerate angular prior rotation.');
  const x = a.map(v => v / na), dot = x.reduce((s, v, i) => s + v * b[i], 0);
  const y = b.map((v, i) => v - dot * x[i]), nb = Math.hypot(...y);
  if (nb < 1e-6) throw new Error('Degenerate angular prior rotation.');
  for (let i = 0; i < 3; i++) y[i] /= nb;
  const z = [x[1]*y[2]-x[2]*y[1], x[2]*y[0]-x[0]*y[2], x[0]*y[1]-x[1]*y[0]];
  return [x, y, z]; // Columns, matching the Python six-dimensional representation.
}

function addAngularPrior(input, output) {
  for (let j = 0; j < 22; j++) {
    const previous = rotation(input, 276 + 6*j), current = rotation(input, 552 + 6*j);
    for (let col = 0; col < 2; col++) {
      // current * transpose(previous) * currentColumn
      const intermediate = previous.map(p => p.reduce((s, v, i) => s + v*current[col][i], 0));
      for (let row = 0; row < 3; row++) {
        const extrapolated = current.reduce((s, c, i) => s + c[row]*intermediate[i], 0);
        output[6*j + 3*col + row] += extrapolated-current[col][row];
      }
    }
  }
}