Determine neighbour indexes for each bubble slot

main.lua

@@ -4,22 +4,66 @@ local game = {}
 local function load_level(level, bubble_diameter, row_gap)
   local slots = {}
   local bubble_radius = bubble_diameter / 2
+  local num_rows = #level
+  local max_cols = math.max(#level[1], #level[2])
+  local min_cols = math.min(#level[1], #level[2])
+
+  assert(
+    max_cols == min_cols + 1,
+    'number of columns in alternating rows must differ by one'
+  )
+
+  local num_bubble_slots = 0
+  if num_rows % 2 == 0 then
+    num_bubble_slots = num_rows / 2 * #level[1] +
+                       num_rows / 2 * #level[2]
+  else
+    num_bubble_slots = math.ceil(num_rows / 2) * #level[1] +
+                       math.floor(num_rows / 2) * #level[2]
+  end
+  print('slots', num_bubble_slots)
 
-  -- set prev_num_cols to size of 2nd row so x_offset is correct when row == 1
-  local prev_num_cols = #level[2]
   for row = 1, #level do
     local num_cols = #level[row]
-    local x_offset = num_cols < prev_num_cols and bubble_radius or 0
+
+    assert(
+      num_cols == min_cols or num_cols == max_cols,
+      string.format("wrong number of columns (%d) in row %d", num_cols, row)
+    )
+
+    local x_offset = num_cols == min_cols and bubble_radius or 0
     for col = 1, num_cols do
-      slots[#slots+1] = {
+      local index = #slots + 1
+      slots[index] = {
         x = x_offset + bubble_radius + (col - 1) * bubble_diameter,
         y = bubble_radius + (row - 1) * row_gap,
-        bubble_type = level[row][col]
+        bubble_type = level[row][col],
+        neighbours = {}
       }
+      if col > 1 then -- left
+        slots[index].neighbours[#slots[index].neighbours+1] = index - 1
+      end
+      if col < num_cols then -- right
+        slots[index].neighbours[#slots[index].neighbours+1] = index + 1
+      end
+      if row % 2 == 0 or (row > 1 and col > 1) then -- up left
+        slots[index].neighbours[#slots[index].neighbours+1] = index - max_cols
+      end
+      if row % 2 == 0 or (row > 1 and col < num_cols) then -- up right
+        slots[index].neighbours[#slots[index].neighbours+1] = index - min_cols
+      end
+      if row % 2 == 0 or col > 1 then -- down left
+        if index + min_cols <= num_bubble_slots then
+          slots[index].neighbours[#slots[index].neighbours+1] = index + min_cols
+        end
+      end
+      if row % 2 == 0 or col < num_cols then -- down right
+        if index + max_cols <= num_bubble_slots then
+          slots[index].neighbours[#slots[index].neighbours+1] = index + max_cols
+        end
+      end
     end
-    prev_num_cols = num_cols
   end
-
   return slots
 end