///|
priv struct WaveSearch {
  model : RuleModel
  width : Int
  height : Int
  periodic : Bool
  count : Int
  domains : FixedArray[Bool]
  supports : FixedArray[Int]
  sizes : FixedArray[Int]
  sums : FixedArray[Double]
  logs : FixedArray[Double]
  weight_logs : Array[Double]
  versions : FixedArray[Int]
  trail : Array[Int]
  processed : Array[Int]
  queue : Array[Int]
  heap : Array[EntropyEntry]
  mut budget : Int
  mut seed : UInt
  mut decisions : Int
  mut backtracks : Int
}

///|
priv struct Decision {
  cell : Int
  trail_mark : Int
  processed_mark : Int
  choices : Array[Int]
}

///|
fn WaveSearch::tick(self : WaveSearch) -> Unit raise SolveError {
  self.budget -= 1
  if self.budget < 0 {
    raise BudgetExhausted
  }
}

///|
fn WaveSearch::random(self : WaveSearch) -> Double {
  let mut x = self.seed
  x = x ^ (x << 13)
  x = x ^ (x >> 17)
  x = x ^ (x << 5)
  self.seed = x
  x.to_double() / 4294967296.0
}

///|
fn WaveSearch::enqueue_entropy(self : WaveSearch, cell : Int) -> Unit {
  self.versions[cell] += 1
  if self.sizes[cell] > 1 {
    let sum = self.sums[cell]
    let entropy = @math.ln(sum) - self.logs[cell] / sum
    heap_push(self.heap, {
      cell,
      version: self.versions[cell],
      priority: entropy + self.random() * 1.0e-9,
    })
  }
}

///|
fn WaveSearch::ban(
  self : WaveSearch,
  cell : Int,
  tile : Int,
) -> Unit raise SolveError {
  let index = cell * self.count + tile
  if !self.domains[index] {
    return
  }
  self.tick()
  self.domains[index] = false
  self.sizes[cell] -= 1
  self.sums[cell] -= self.model.weights[tile]
  self.logs[cell] -= self.weight_logs[tile]
  // Removing a dominant weight can erase tiny remaining weights by rounding.
  // Rebuild the two sums before computing entropy in that case.
  if self.sizes[cell] > 1 &&
    self.sums[cell] <= self.model.weights[tile] * 1.0e-8 {
    self.sums[cell] = 0.0
    self.logs[cell] = 0.0
    for other in 0.. self.sizes.length() * 8 + 1024 {
    self.heap.clear()
    for cell in 0.. Bool raise SolveError {
  let mut valid = true
  while self.queue.pop() is Some(index) {
    self.tick()
    let cell = index / self.count
    let tile = index % self.count
    if self.sizes[cell] == 0 {
      valid = false
    }
    // Finish the event before returning a contradiction so rollback can restore
    // exactly the counters changed by each processed removal.
    for direction in 0..<4 {
      if grid_neighbor(cell, direction, self.width, self.height, self.periodic)
        is Some(other) {
        for target in self.model.neighbors[tile][direction] {
          self.tick()
          let counter = (other * self.count + target) * 4 + (direction + 2) % 4
          self.supports[counter] -= 1
          if self.supports[counter] == 0 {
            self.ban(other, target)
          }
        }
      }
    }
    self.processed.push(index)
  }
  valid
}

///|
fn WaveSearch::rollback(
  self : WaveSearch,
  decision : Decision,
) -> Unit raise SolveError {
  self.queue.clear()
  while self.processed.length() > decision.processed_mark {
    let index = self.processed.pop().unwrap()
    let cell = index / self.count
    let tile = index % self.count
    for direction in 0..<4 {
      if grid_neighbor(cell, direction, self.width, self.height, self.periodic)
        is Some(other) {
        for target in self.model.neighbors[tile][direction] {
          self.tick()
          self.supports[(other * self.count + target) * 4 + (direction + 2) % 4] += 1
        }
      }
    }
  }
  let touched : Map[Int, Bool] = Map([])
  while self.trail.length() > decision.trail_mark {
    let index = self.trail.pop().unwrap()
    let cell = index / self.count
    let tile = index % self.count
    self.domains[index] = true
    self.sizes[cell] += 1
    self.sums[cell] += self.model.weights[tile]
    self.logs[cell] += self.weight_logs[tile]
    touched[cell] = true
  }
  for cell, _ in touched {
    self.enqueue_entropy(cell)
  }
}

///|
fn WaveSearch::collapse(
  self : WaveSearch,
  decision : Decision,
) -> Unit raise SolveError {
  let mut total = 0.0
  for tile in decision.choices {
    total += self.model.weights[tile]
  }
  let mut draw = self.random() * total
  let mut picked = decision.choices.length() - 1
  for i, tile in decision.choices {
    draw -= self.model.weights[tile]
    if draw < 0.0 {
      picked = i
      break
    }
  }
  let tile = decision.choices[picked]
  decision.choices[picked] = decision.choices[decision.choices.length() - 1]
  ignore(decision.choices.pop())
  for other in 0.. Solution? raise SolveError {
  model.checked()
  let count = model.labels.length()
  if width < 1 ||
    height < 1 ||
    width > 65536 ||
    height > 65536 ||
    width.to_int64() * height.to_int64() > 65536L ||
    width.to_int64() * height.to_int64() * count.to_int64() > 2000000L ||
    budget < 1 {
    raise Invalid("grid/state dimensions or budget exceed limits")
  }
  let cells = width * height
  let states = cells * count
  if pins.length() > 65536 || restrictions.length() > 65536 {
    raise Invalid("constraint count limit")
  }
  // Normalize repeated pins before wave allocation. Identical pins add no
  // constraint; conflicting pins become one empty domain. Preserve the first
  // occurrence order so distinct input pins keep their existing ordering.
  let pinned_tiles : Map[Int, Int] = Map([])
  let pinned_cells : Array[Int] = []
  for (cell, tile) in pins {
    if cell < 0 || cell >= cells || tile < 0 || tile >= count {
      raise Invalid("invalid pin")
    }
    match pinned_tiles.get(cell) {
      None => {
        pinned_tiles[cell] = tile
        pinned_cells.push(cell)
      }
      Some(previous) => if previous != tile { pinned_tiles[cell] = -1 }
    }
  }
  // Intersect repeated cell restrictions before touching the wave. Previously
  // each repeated restriction scanned every tile even after all bans were
  // already applied, so that work did not consume the propagation budget.
  let cell_rules : Map[Int, Map[Int, Bool]] = Map([])
  let mut candidates = 0
  for (cell, allowed) in restrictions {
    if cell < 0 || cell >= cells || allowed.length() > 2000000 - candidates {
      raise Invalid("invalid cell restriction or candidate limit")
    }
    candidates += allowed.length()
    let keep : Map[Int, Bool] = Map([])
    for tile in allowed {
      if tile < 0 || tile >= count {
        raise Invalid("invalid cell restriction")
      }
      keep[tile] = true
    }
    match cell_rules.get(cell) {
      None => cell_rules[cell] = keep
      Some(previous) =>
        for tile in previous.keys().to_array() {
          if !keep.contains(tile) {
            previous.remove(tile)
          }
        }
    }
  }
  let mut total = 0.0
  let weight_logs = model.weights.map(w => w * @math.ln(w))
  let mut total_log = 0.0
  for i, weight in model.weights {
    total += weight
    total_log += weight_logs[i]
  }
  let self : WaveSearch = {
    model,
    width,
    height,
    periodic,
    count,
    domains: FixedArray::make(states, true),
    supports: FixedArray::make(states * 4, 0),
    sizes: FixedArray::make(cells, count),
    sums: FixedArray::make(cells, total),
    logs: FixedArray::make(cells, total_log),
    weight_logs,
    versions: FixedArray::make(cells, 0),
    trail: [],
    processed: [],
    queue: [],
    heap: [],
    budget,
    seed: if seed == 0U {
      1U
    } else {
      seed
    },
    decisions: 0,
    backtracks: 0,
  }
  for cell in 0.. 1 {
        selected = entry.cell
        break
      }
    }
    if selected < 0 {
      let tiles = Array::makei(cells, cell => {
        let mut tile = 0
        while tile < count && !self.domains[cell * count + tile] {
          tile += 1
        }
        tile
      })
      return Some({
        width,
        height,
        tiles,
        decisions: self.decisions,
        backtracks: self.backtracks,
      })
    }
    let decision : Decision = {
      cell: selected,
      trail_mark: self.trail.length(),
      processed_mark: self.processed.length(),
      choices: Array::makei(count, i => i).filter(t => {
        self.domains[selected * count + t]
      }),
    }
    stack.push(decision)
    self.decisions += 1
    self.collapse(decision)
    while !self.propagate() {
      let mut resumed = false
      while stack.pop() is Some(previous) {
        self.rollback(previous)
        self.backtracks += 1
        if !previous.choices.is_empty() {
          stack.push(previous)
          self.collapse(previous)
          resumed = true
          break
        }
      }
      if !resumed {
        return None
      }
    }
  }
  None
}