/**
 * Two-pass frequency-domain delay-and-sum beamforming kernel.
 *
 * Dispatch shapes:
 *   compute_steering = (ceil(microphoneCount / WORKGROUP_SIZE), pointCount, binCount)
 *   compute_das      = (ceil(pointCount / WORKGROUP_SIZE), binCount, 1)
 *
 * The first pass computes trigonometry once per (bin, point, microphone).
 * The second pass reuses those steering vectors for every STFT frame. Callers
 * can bound the intermediate buffer by submitting point/bin batches.
 *
 * Flat buffer layouts:
 *   spectrum:    [bin][frame][microphone][real, imaginary]
 *   distances:   [point][microphone]
 *   calibration: [microphone][inverse sensitivity gain or unity, phase offset radians]
 *   steering:    [bin][point][microphone][real, imaginary]
 *   output:      [bin][point]
 */

export const WORKGROUP_SIZE = 64;
export const DAS_WORKGROUP_SIZE = WORKGROUP_SIZE;
export const MAX_MICROPHONES = 1024;

// WGSL uniform structs are rounded to a 16-byte boundary. The reserved word
// keeps the JavaScript and shader representations explicitly 32 bytes.
export const DAS_PARAMS_BYTE_LENGTH = 32;

export const DAS_PARAMS_LAYOUT = Object.freeze({
  byteLength: DAS_PARAMS_BYTE_LENGTH,
  pointCount: Object.freeze({ offset: 0, type: "u32" }),
  microphoneCount: Object.freeze({ offset: 4, type: "u32" }),
  frameCount: Object.freeze({ offset: 8, type: "u32" }),
  binCount: Object.freeze({ offset: 12, type: "u32" }),
  sampleRate: Object.freeze({ offset: 16, type: "f32" }),
  fftSize: Object.freeze({ offset: 20, type: "u32" }),
  speedOfSound: Object.freeze({ offset: 24, type: "f32" }),
  reserved: Object.freeze({ offset: 28, type: "u32" }),
});

export const DAS_BUFFER_LAYOUT = Object.freeze({
  params: Object.freeze({ group: 0, binding: 0, access: "uniform", layout: DAS_PARAMS_LAYOUT }),
  binIndices: Object.freeze({ group: 0, binding: 1, access: "read", element: "u32" }),
  spectrum: Object.freeze({
    group: 0,
    binding: 2,
    access: "read",
    element: "f32",
    order: "bin-frame-microphone-complex",
  }),
  distances: Object.freeze({
    group: 0,
    binding: 3,
    access: "read",
    element: "f32",
    order: "point-microphone",
  }),
  calibration: Object.freeze({
    group: 0,
    binding: 4,
    access: "read",
    element: "vec2<f32>",
    order: "microphone-(inverseGain-phaseRadians)",
  }),
  steering: Object.freeze({
    group: 0,
    binding: 5,
    access: "read_write",
    element: "vec2<f32>",
    order: "bin-point-microphone-complex",
  }),
  output: Object.freeze({
    group: 0,
    binding: 6,
    access: "read_write",
    element: "f32",
    order: "bin-point",
  }),
});

export const DAS_ENTRY_POINTS = Object.freeze({
  steering: "compute_steering",
  beamforming: "compute_das",
});

function requireUint32(name, value) {
  if (!Number.isInteger(value) || value < 0 || value > 0xffff_ffff) {
    throw new RangeError(`${name} must be an unsigned 32-bit integer.`);
  }
  return value;
}

function requireUint32Product(name, values) {
  let product = 1;
  for (const value of values) {
    if (value !== 0 && product > 0xffff_ffff / value) {
      throw new RangeError(`${name} exceeds the shader's u32 address space.`);
    }
    product *= value;
  }
  return product;
}

function requirePositiveF32(name, value) {
  const rounded = Math.fround(value);
  if (!(rounded > 0) || !Number.isFinite(rounded)) {
    throw new RangeError(`${name} must be a finite positive f32 value.`);
  }
  return rounded;
}

function requireFiniteF32(name, value) {
  const rounded = Math.fround(value);
  if (!Number.isFinite(rounded)) {
    throw new RangeError(`${name} must be a finite f32 value.`);
  }
  return rounded;
}

function requireFiniteNumber(name, value) {
  if (!Number.isFinite(value)) throw new RangeError(`${name} must be finite.`);
  return value;
}

/** Pack the binding-0 uniform using the offsets in DAS_PARAMS_LAYOUT. */
export function packDasParams(options) {
  const {
    pointCount,
    microphoneCount,
    frameCount,
    binCount,
    sampleRate,
    fftSize,
    speedOfSound,
  } = options;
  const counts = {
    pointCount: requireUint32("pointCount", pointCount),
    microphoneCount: requireUint32("microphoneCount", microphoneCount),
    frameCount: requireUint32("frameCount", frameCount),
    binCount: requireUint32("binCount", binCount),
    fftSize: requireUint32("fftSize", fftSize),
  };
  if (counts.microphoneCount > MAX_MICROPHONES) {
    throw new RangeError(`microphoneCount must not exceed ${MAX_MICROPHONES}.`);
  }
  if (counts.fftSize === 0) throw new RangeError("fftSize must be greater than zero.");

  requireUint32Product("distance buffer element count", [
    counts.pointCount,
    counts.microphoneCount,
  ]);
  requireUint32Product("output buffer element count", [counts.binCount, counts.pointCount]);
  requireUint32Product("steering buffer element count", [
    counts.binCount,
    counts.pointCount,
    counts.microphoneCount,
  ]);
  requireUint32Product("spectrum buffer element count", [
    counts.binCount,
    counts.frameCount,
    counts.microphoneCount,
    2,
  ]);

  const packedSampleRate = requirePositiveF32("sampleRate", sampleRate);
  const packedSpeedOfSound = requirePositiveF32("speedOfSound", speedOfSound);
  const buffer = new ArrayBuffer(DAS_PARAMS_BYTE_LENGTH);
  const view = new DataView(buffer);
  view.setUint32(0, counts.pointCount, true);
  view.setUint32(4, counts.microphoneCount, true);
  view.setUint32(8, counts.frameCount, true);
  view.setUint32(12, counts.binCount, true);
  view.setFloat32(16, packedSampleRate, true);
  view.setUint32(20, counts.fftSize, true);
  view.setFloat32(24, packedSpeedOfSound, true);
  view.setUint32(28, 0, true);
  return buffer;
}

/**
 * Pack already-linear calibration as tightly packed vec2<f32> entries.
 * phase is the microphone phase offset in radians. The shader subtracts it
 * from the propagation phase, exactly as the CPU reference does.
 */
export function packDasCalibration(calibration) {
  if (!Array.isArray(calibration)) throw new TypeError("calibration must be an array.");
  const packed = new Float32Array(calibration.length * 2);
  for (let index = 0; index < calibration.length; index += 1) {
    const entry = calibration[index];
    const gain = requireFiniteF32(`calibration[${index}].gain`, entry?.gain);
    if (gain < 0) throw new RangeError(`calibration[${index}].gain must not be negative.`);
    packed[index * 2] = gain;
    packed[index * 2 + 1] = requireFiniteF32(
      `calibration[${index}].phase`,
      entry?.phase,
    );
  }
  return packed;
}

/**
 * Convert array-definition microphone metadata to the shader calibration layout.
 * Phase-only algorithms set includeAmplitude=false while retaining phase offsets.
 */
export function packMicrophoneCalibration(microphones, options = {}) {
  if (!Array.isArray(microphones)) throw new TypeError("microphones must be an array.");
  const includeAmplitude = options.includeAmplitude !== false;
  const packed = new Float32Array(microphones.length * 2);
  for (let index = 0; index < microphones.length; index += 1) {
    const microphone = microphones[index];
    const sensitivityDb = requireFiniteNumber(
      `microphones[${index}].sensitivityDb`,
      microphone?.sensitivityDb ?? 0,
    );
    const phaseOffsetDeg = requireFiniteNumber(
      `microphones[${index}].phaseOffsetDeg`,
      microphone?.phaseOffsetDeg ?? 0,
    );
    packed[index * 2] = requireFiniteF32(
      `microphones[${index}] inverse sensitivity gain`,
      includeAmplitude ? 10 ** (-sensitivityDb / 20) : 1,
    );
    packed[index * 2 + 1] = requireFiniteF32(
      `microphones[${index}] phase offset`,
      (phaseOffsetDeg * Math.PI) / 180,
    );
  }
  return packed;
}

export const DAS_WGSL = /* wgsl */ `
struct DasParams {
  point_count: u32,
  microphone_count: u32,
  frame_count: u32,
  bin_count: u32,
  sample_rate: f32,
  fft_size: u32,
  speed_of_sound: f32,
  reserved: u32,
}

struct U32Buffer {
  values: array<u32>,
}

struct F32Buffer {
  values: array<f32>,
}

struct Vec2Buffer {
  values: array<vec2<f32>>,
}

@group(0) @binding(0) var<uniform> params: DasParams;
@group(0) @binding(1) var<storage, read> bin_indices: U32Buffer;
@group(0) @binding(2) var<storage, read> spectrum: F32Buffer;
@group(0) @binding(3) var<storage, read> distances: F32Buffer;
@group(0) @binding(4) var<storage, read> calibration: Vec2Buffer;
@group(0) @binding(5) var<storage, read_write> steering_vectors: Vec2Buffer;
@group(0) @binding(6) var<storage, read_write> output_maps: F32Buffer;

const TAU: f32 = 6.283185307179586;
const MAX_MICROPHONES: u32 = ${MAX_MICROPHONES}u;
const MAX_U32: u32 = 0xffffffffu;

fn multiplication_overflows(left: u32, right: u32) -> bool {
  return left != 0u && right > MAX_U32 / left;
}

fn wrap_phase(phase: f32) -> f32 {
  return phase - round(phase / TAU) * TAU;
}

@compute @workgroup_size(${WORKGROUP_SIZE}, 1, 1)
fn compute_steering(@builtin(global_invocation_id) invocation: vec3<u32>) {
  let microphone_index = invocation.x;
  let point_index = invocation.y;
  let bin_index = invocation.z;

  if (
    microphone_index >= params.microphone_count ||
    point_index >= params.point_count ||
    bin_index >= params.bin_count
  ) {
    return;
  }

  if (
    params.microphone_count == 0u ||
    params.microphone_count > MAX_MICROPHONES ||
    params.fft_size == 0u ||
    params.sample_rate <= 0.0 ||
    params.speed_of_sound <= 0.0 ||
    multiplication_overflows(params.point_count, params.microphone_count) ||
    multiplication_overflows(params.bin_count, params.point_count)
  ) {
    return;
  }

  let bin_point_count = params.bin_count * params.point_count;
  if (multiplication_overflows(bin_point_count, params.microphone_count)) {
    return;
  }

  let required_distance_values = params.point_count * params.microphone_count;
  let required_steering_values = bin_point_count * params.microphone_count;
  if (
    arrayLength(&bin_indices.values) < params.bin_count ||
    arrayLength(&calibration.values) < params.microphone_count ||
    arrayLength(&distances.values) < required_distance_values ||
    arrayLength(&steering_vectors.values) < required_steering_values
  ) {
    return;
  }

  let bin = f32(bin_indices.values[bin_index]);
  let angular_wave_number =
    (TAU * bin * params.sample_rate) /
    (f32(params.fft_size) * params.speed_of_sound);
  let distance_index = point_index * params.microphone_count + microphone_index;
  let steering_index =
    (bin_index * params.point_count + point_index) * params.microphone_count +
    microphone_index;
  let microphone_calibration = calibration.values[microphone_index];
  let phase = wrap_phase(
    angular_wave_number * distances.values[distance_index] -
    microphone_calibration.y
  );
  steering_vectors.values[steering_index] =
    vec2<f32>(cos(phase), sin(phase)) * microphone_calibration.x;
}

@compute @workgroup_size(${WORKGROUP_SIZE}, 1, 1)
fn compute_das(@builtin(global_invocation_id) invocation: vec3<u32>) {
  let point_index = invocation.x;
  let bin_index = invocation.y;

  if (point_index >= params.point_count || bin_index >= params.bin_count) {
    return;
  }

  if (
    params.microphone_count == 0u ||
    params.microphone_count > MAX_MICROPHONES ||
    params.frame_count == 0u ||
    multiplication_overflows(params.bin_count, params.point_count) ||
    multiplication_overflows(params.bin_count, params.frame_count)
  ) {
    return;
  }

  let bin_point_count = params.bin_count * params.point_count;
  let bin_frame_count = params.bin_count * params.frame_count;
  if (
    multiplication_overflows(bin_point_count, params.microphone_count) ||
    multiplication_overflows(bin_frame_count, params.microphone_count)
  ) {
    return;
  }

  let spectrum_complex_count = bin_frame_count * params.microphone_count;
  if (multiplication_overflows(spectrum_complex_count, 2u)) {
    return;
  }

  let required_output_values = bin_point_count;
  let required_steering_values = bin_point_count * params.microphone_count;
  let required_spectrum_values = spectrum_complex_count * 2u;
  if (
    arrayLength(&output_maps.values) < required_output_values ||
    arrayLength(&steering_vectors.values) < required_steering_values ||
    arrayLength(&spectrum.values) < required_spectrum_values
  ) {
    return;
  }

  let microphone_scale = 1.0 / f32(params.microphone_count);
  let steering_base =
    (bin_index * params.point_count + point_index) * params.microphone_count;
  var accumulated_power = 0.0;

  for (var frame_index = 0u; frame_index < params.frame_count; frame_index += 1u) {
    var sum = vec2<f32>(0.0, 0.0);
    for (
      var microphone_index = 0u;
      microphone_index < params.microphone_count;
      microphone_index += 1u
    ) {
      let complex_index = (
        (
          (bin_index * params.frame_count + frame_index) * params.microphone_count +
          microphone_index
        ) * 2u
      );
      let input = vec2<f32>(
        spectrum.values[complex_index],
        spectrum.values[complex_index + 1u]
      );
      let steering = steering_vectors.values[steering_base + microphone_index];

      // (input.real + i*input.imaginary) *
      // (steering.real + i*steering.imaginary)
      sum += vec2<f32>(
        input.x * steering.x - input.y * steering.y,
        input.x * steering.y + input.y * steering.x
      );
    }

    accumulated_power +=
      dot(sum, sum) * microphone_scale * microphone_scale;
  }

  let output_index = bin_index * params.point_count + point_index;
  output_maps.values[output_index] = accumulated_power / f32(params.frame_count);
}
`;

export const DAS_SHADER_SOURCE = DAS_WGSL;