// Poisson Disc Sampling — GPU-accelerated Bridson's algorithm
//
// Generates 2D points with a guaranteed minimum distance between them.
// Iterates until no more active points can spawn valid candidates.
//
// Single kernel (generate) dispatched repeatedly from the host.
// The host swaps active list regions between iterations.

struct Uniforms {
  width: f32,
  height: f32,
  min_distance: f32,
  cell_size: f32,       // min_distance / sqrt(2)
  grid_width: u32,
  grid_height: u32,
  seed: u32,
  max_points: u32,
  num_candidates: u32,  // candidates per active_list point
  active_list_offset: u32,   // read offset into active_list buffer (ping-pong)
  write_offset: u32,    // write offset into active_list buffer (ping-pong)
  _pad: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read_write> points: array<vec2f>;
@group(0) @binding(2) var<storage, read_write> grid: array<atomic<i32>>;
@group(0) @binding(3) var<storage, read_write> active_list: array<u32>;
@group(0) @binding(4) var<storage, read_write> counters: array<atomic<u32>>;
// counters[0] = point count
// counters[1] = active_list count (current iteration, read-only during dispatch)
// counters[2] = next active_list count (written during dispatch)

// Sentinel value: cell is claimed but point data not yet readable.
// is_valid() treats this as occupied (rejects candidates near it).
const CELL_EMPTY: i32 = -1i;
const CELL_PENDING: i32 = -2i;

// --- Hash-based RNG ---
// PCG-style hash for generating random candidates

fn pcg(v: u32) -> u32 {
  var state = v * 747796405u + 2891336453u;
  let word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
  return (word >> 22u) ^ word;
}

fn rand_f32(seed: u32) -> f32 {
  return f32(pcg(seed)) / 4294967295.0;
}

// Generate a random point in the annulus [minDistance, 2*minDistance] around center
fn random_in_annulus(center: vec2f, seed: u32) -> vec2f {
  let r1 = rand_f32(seed);
  let r2 = rand_f32(seed + 1u);
  let radius = u.min_distance * (1.0 + r1);
  let angle = r2 * 6.283185;
  return center + vec2f(cos(angle), sin(angle)) * radius;
}

// Convert a point to a grid cell index
fn point_to_cell(p: vec2f) -> vec2u {
  return vec2u(
    u32(p.x / u.cell_size),
    u32(p.y / u.cell_size)
  );
}

fn cell_to_index(cell: vec2u) -> u32 {
  return cell.y * u.grid_width + cell.x;
}

// Check if a point is far enough from all neighbours in the grid
fn is_valid(p: vec2f) -> bool {
  if (p.x < 0.0 || p.x >= u.width || p.y < 0.0 || p.y >= u.height) {
    return false;
  }

  let cell = point_to_cell(p);
  let min_dist_sq = u.min_distance * u.min_distance;

  // Check 5x5 neighbourhood of cells
  for (var dy = -2i; dy <= 2i; dy++) {
    for (var dx = -2i; dx <= 2i; dx++) {
      let nx = i32(cell.x) + dx;
      let ny = i32(cell.y) + dy;

      if (nx < 0i || nx >= i32(u.grid_width) || ny < 0i || ny >= i32(u.grid_height)) {
        continue;
      }

      let neighbour_idx = atomicLoad(&grid[u32(ny) * u.grid_width + u32(nx)]);

      // Empty cell — no point here
      if (neighbour_idx == CELL_EMPTY) {
        continue;
      }

      // Pending cell — another thread claimed it but hasn't written yet.
      // Treat as occupied to be safe (reject candidate).
      if (neighbour_idx == CELL_PENDING) {
        return false;
      }

      let neighbour = points[u32(neighbour_idx)];
      let d = p - neighbour;
      if (dot(d, d) < min_dist_sq) {
        return false;
      }
    }
  }

  return true;
}

// ============================================================
// Kernel: Generate candidates and accept valid ones
// Each thread processes one active_list point, tries multiple candidates.
// Active list uses ping-pong offsets to avoid read/write aliasing.
// ============================================================

@compute @workgroup_size(64)
fn generate(@builtin(global_invocation_id) gid: vec3u) {
  let active_list_count = atomicLoad(&counters[1]);
  if (gid.x >= active_list_count) {
    return;
  }

  let point_idx = active_list[u.active_list_offset + gid.x];
  let center = points[point_idx];
  var found = false;

  for (var k = 0u; k < u.num_candidates; k++) {
    let hash_seed = u.seed + point_idx * 31u + k * 7u + gid.x * 113u;
    let candidate = random_in_annulus(center, hash_seed);

    if (!is_valid(candidate)) {
      continue;
    }

    // Try to claim the grid cell with an atomic compare-exchange
    let cell = point_to_cell(candidate);
    let cell_idx = cell_to_index(cell);

    // Set to CELL_PENDING first — is_valid() will treat this as occupied
    let result = atomicCompareExchangeWeak(&grid[cell_idx], CELL_EMPTY, CELL_PENDING);
    if (!result.exchanged) {
      continue; // Cell already taken
    }

    // Claim a point slot
    let new_idx = atomicAdd(&counters[0], 1u);
    if (new_idx >= u.max_points) {
      atomicSub(&counters[0], 1u);
      atomicStore(&grid[cell_idx], CELL_EMPTY);
      return;
    }

    // Write the point, then update grid to the real index
    points[new_idx] = candidate;
    atomicStore(&grid[cell_idx], i32(new_idx));

    // Add to next iteration's active list (write region)
    let next_active_list_idx = atomicAdd(&counters[2], 1u);
    if (next_active_list_idx < u.max_points) {
      active_list[u.write_offset + next_active_list_idx] = new_idx;
    }

    found = true;
  }

  // If this active_list point generated at least one child, keep it active_list
  if (found) {
    let keep_idx = atomicAdd(&counters[2], 1u);
    if (keep_idx < u.max_points) {
      active_list[u.write_offset + keep_idx] = point_idx;
    }
  }
}

// Poisson Disc Sampling
// Generates 2D points with a guaranteed minimum distance between them.
// One-shot async operation — call it, get the result, use the buffer.
//
// Default WGSL loading uses a ?raw import (works with Vite, esbuild, Webpack).
// Alternative: load via fetch — see README.md for details.
import shaderSource from './poisson-disc.wgsl?raw';

export interface PoissonDiscOptions {
  width?: number;
  height?: number;
  minDistance?: number;
  seed?: number;
  maxPoints?: number;
  candidatesPerPoint?: number;
  maxIterations?: number;
  /** Run a CPU post-filter to remove points that violate the minimum distance.
   *  The GPU's parallel algorithm can produce rare violations due to race
   *  conditions between workgroups. Default false. */
  strict?: boolean;
}

export interface PoissonDiscResult {
  points: GPUBuffer;
  count: number;
  destroy(): void;
}

const WORKGROUP_SIZE = 64;

// Uniform buffer layout (48 bytes, 12 x u32/f32):
// f32 width, f32 height, f32 min_distance, f32 cell_size,
// u32 grid_width, u32 grid_height, u32 seed, u32 max_points,
// u32 num_candidates, u32 active_offset, u32 write_offset, u32 _pad
const UNIFORM_SIZE = 48;

export async function poissonDisc(
  device: GPUDevice,
  options: PoissonDiscOptions = {}
): Promise<PoissonDiscResult> {
  const width = options.width ?? 100;
  const height = options.height ?? 100;
  const minDistance = options.minDistance ?? 2.0;
  const seed = options.seed ?? Math.floor(Math.random() * 0xffffffff);
  const maxPoints = options.maxPoints ?? 10_000;
  const candidatesPerPoint = options.candidatesPerPoint ?? 30;
  const maxIterations = options.maxIterations ?? 500;
  const strict = options.strict ?? false;

  const cellSize = minDistance / Math.SQRT2;
  const gridWidth = Math.ceil(width / cellSize);
  const gridHeight = Math.ceil(height / cellSize);

  // Active buffer uses ping-pong: two regions of maxPoints each
  const activeListBufferSize = maxPoints * 2;

  // Buffers
  const pointsBuffer = device.createBuffer({
    size: maxPoints * 8, // vec2<f32> = 8 bytes
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
  });

  const gridBuffer = device.createBuffer({
    size: gridWidth * gridHeight * 4, // i32 per cell
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
  });

  const activeListBuffer = device.createBuffer({
    size: activeListBufferSize * 4, // u32 indices, two regions
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
  });

  // counters[0] = point count, counters[1] = active count, counters[2] = next active count
  const countersBuffer = device.createBuffer({
    size: 12,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
  });

  const readbackBuffer = device.createBuffer({
    size: 12,
    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
  });

  const uniformBuffer = device.createBuffer({
    size: UNIFORM_SIZE,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
  });

  // Initialize grid to -1 (empty)
  const gridInit = new Int32Array(gridWidth * gridHeight).fill(-1);
  device.queue.writeBuffer(gridBuffer, 0, gridInit);

  // Place initial seed point at the center
  const seedPoint = new Float32Array([width / 2, height / 2]);
  device.queue.writeBuffer(pointsBuffer, 0, seedPoint);

  const seedCellX = Math.floor(width / 2 / cellSize);
  const seedCellY = Math.floor(height / 2 / cellSize);
  const seedCellIdx = seedCellY * gridWidth + seedCellX;
  device.queue.writeBuffer(gridBuffer, seedCellIdx * 4, new Int32Array([0]));

  // Initialize counters: 1 point, 1 active, 0 next active
  device.queue.writeBuffer(countersBuffer, 0, new Uint32Array([1, 1, 0]));

  // Initialize active list with seed point (index 0) in region A (offset 0)
  device.queue.writeBuffer(activeListBuffer, 0, new Uint32Array([0]));

  // Shader and pipeline
  const shaderModule = device.createShaderModule({ code: shaderSource });
  const compilationInfo = await shaderModule.getCompilationInfo();
  const errors = compilationInfo.messages.filter((m) => m.type === 'error');
  if (errors.length > 0) {
    const details = errors.map((m) => `line ${m.lineNum}: ${m.message}`).join('\n');
    throw new Error(`[poisson-disc] WGSL compilation failed:\n${details}`);
  }

  const bindGroupLayout = device.createBindGroupLayout({
    entries: [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } }
    ]
  });

  const pipeline = device.createComputePipeline({
    layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
    compute: { module: shaderModule, entryPoint: 'generate' }
  });

  const bindGroup = device.createBindGroup({
    layout: bindGroupLayout,
    entries: [
      { binding: 0, resource: { buffer: uniformBuffer } },
      { binding: 1, resource: { buffer: pointsBuffer } },
      { binding: 2, resource: { buffer: gridBuffer } },
      { binding: 3, resource: { buffer: activeListBuffer } },
      { binding: 4, resource: { buffer: countersBuffer } }
    ]
  });

  // Uniform data (reused each iteration, only offsets change)
  const uniformData = new ArrayBuffer(UNIFORM_SIZE);
  const f32 = new Float32Array(uniformData);
  const u32 = new Uint32Array(uniformData);
  f32[0] = width;
  f32[1] = height;
  f32[2] = minDistance;
  f32[3] = cellSize;
  u32[4] = gridWidth;
  u32[5] = gridHeight;
  u32[6] = seed;
  u32[7] = maxPoints;
  u32[8] = candidatesPerPoint;

  // Ping-pong: region A = offset 0, region B = offset maxPoints
  let readOffset = 0;
  let writeOffset = maxPoints;

  let activeCount = 1; // seed point

  for (let iter = 0; iter < maxIterations; iter++) {
    if (activeCount === 0) break;

    // Update uniforms with current ping-pong offsets
    u32[9] = readOffset;
    u32[10] = writeOffset;
    u32[11] = 0; // padding
    device.queue.writeBuffer(uniformBuffer, 0, uniformData);

    // Set active count and reset next active count on the GPU
    device.queue.writeBuffer(countersBuffer, 4, new Uint32Array([activeCount]));
    device.queue.writeBuffer(countersBuffer, 8, new Uint32Array([0]));

    // Dispatch generate kernel
    const dispatchCount = Math.ceil(activeCount / WORKGROUP_SIZE);
    const computeEncoder = device.createCommandEncoder();
    const pass = computeEncoder.beginComputePass();
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    pass.dispatchWorkgroups(dispatchCount);
    pass.end();

    // Read back next active count after dispatch
    computeEncoder.copyBufferToBuffer(countersBuffer, 8, readbackBuffer, 0, 4);
    device.queue.submit([computeEncoder.finish()]);

    await readbackBuffer.mapAsync(GPUMapMode.READ);
    activeCount = new Uint32Array(readbackBuffer.getMappedRange().slice(0))[0];
    readbackBuffer.unmap();

    // Swap ping-pong regions
    const tmp = readOffset;
    readOffset = writeOffset;
    writeOffset = tmp;
  }

  // Read final point count
  const countEncoder = device.createCommandEncoder();
  countEncoder.copyBufferToBuffer(countersBuffer, 0, readbackBuffer, 0, 4);
  device.queue.submit([countEncoder.finish()]);

  await readbackBuffer.mapAsync(GPUMapMode.READ);
  const rawCount = new Uint32Array(readbackBuffer.getMappedRange().slice(0))[0];
  readbackBuffer.unmap();

  let finalCount = rawCount;

  if (strict && rawCount > 0) {
    // Post-filter: the GPU's parallel cell claiming can produce points that
    // violate the minimum distance (race between threads in different
    // workgroups). Read back all points and remove violations.
    const pointsReadback = device.createBuffer({
      size: rawCount * 8,
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
    });

    const readEncoder = device.createCommandEncoder();
    readEncoder.copyBufferToBuffer(pointsBuffer, 0, pointsReadback, 0, rawCount * 8);
    device.queue.submit([readEncoder.finish()]);

    await pointsReadback.mapAsync(GPUMapMode.READ);
    const rawPoints = new Float32Array(pointsReadback.getMappedRange().slice(0));
    pointsReadback.unmap();
    pointsReadback.destroy();

    const filterCellSize = minDistance / Math.SQRT2;
    const filterGridW = Math.ceil(width / filterCellSize);
    const filterGridH = Math.ceil(height / filterCellSize);
    const filterGrid = new Int32Array(filterGridW * filterGridH).fill(-1);
    const cleanPoints = new Float32Array(rawCount * 2);
    let cleanCount = 0;
    const minDistSq = minDistance * minDistance;

    for (let i = 0; i < rawCount; i++) {
      const px = rawPoints[i * 2];
      const py = rawPoints[i * 2 + 1];
      const cx = Math.floor(px / filterCellSize);
      const cy = Math.floor(py / filterCellSize);
      let valid = true;

      outer: for (let dy = -2; dy <= 2; dy++) {
        for (let dx = -2; dx <= 2; dx++) {
          const nx = cx + dx;
          const ny = cy + dy;
          if (nx < 0 || nx >= filterGridW || ny < 0 || ny >= filterGridH) continue;
          const idx = filterGrid[ny * filterGridW + nx];
          if (idx < 0) continue;
          const ddx = cleanPoints[idx * 2] - px;
          const ddy = cleanPoints[idx * 2 + 1] - py;
          if (ddx * ddx + ddy * ddy < minDistSq) {
            valid = false;
            break outer;
          }
        }
      }

      if (valid) {
        cleanPoints[cleanCount * 2] = px;
        cleanPoints[cleanCount * 2 + 1] = py;
        if (cx >= 0 && cx < filterGridW && cy >= 0 && cy < filterGridH) {
          filterGrid[cy * filterGridW + cx] = cleanCount;
        }
        cleanCount++;
      }
    }

    device.queue.writeBuffer(pointsBuffer, 0, cleanPoints.subarray(0, cleanCount * 2));
    finalCount = cleanCount;
  }

  function destroy(): void {
    gridBuffer.destroy();
    activeListBuffer.destroy();
    countersBuffer.destroy();
    readbackBuffer.destroy();
    uniformBuffer.destroy();
  }

  return { points: pointsBuffer, count: finalCount, destroy };
}