local ruleset = {
  {2},
  {1,3},
  {2}
}

local offsets = {
  up = function(grid, cell)
    local x,y = cell.coordinates[1],cell.coordinates[2]
    return grid[x][y-1]
  end,
  down = function(grid, cell)
    local x,y = cell.coordinates[1],cell.coordinates[2]
    return grid[x][y+1]
  end,
  left = function(grid, cell)
    local x,y = cell.coordinates[1],cell.coordinates[2]
    if not grid[x-1] then return end
    return grid[x-1][y]
  end,
  right = function(grid, cell)
    local x,y = cell.coordinates[1],cell.coordinates[2]
    if not grid[x+1] then return end
    return grid[x+1][y]
  end,
}

function getAllModules()
  local modules = {}
  
  for i = 1,#ruleset do
    table.insert(modules, i)
  end
  
  return modules
end

function getNeighbours(grid, cell)
  local neighbours = {}
  
  for _,method in pairs(offsets) do
    local neighbour = method(grid, cell)
    
    if #neighbour.modules > 1 then
      table.insert(neighbours, neighbour)
    end
  end
  
  return neighbours
end

function getAvailableCells(grid)
  local cells = {}
  for i = 1,#grid do
    for j = 1,#grid[1] do
      local cell = grid[i][j]
      
      if #cell.modules == 1 then
        local neighbours = getNeighbours(grid, cell)
        
        if #neighbours > 0 then
          cells[cell] = neighbours
        end
      end
    end
  end
  
  return cells
end

function createGrid(height, width)
  local grid = {}
  for i = 1,height do
    grid[i] = {}
    
    for j = 1,width do
      grid[i][j] = {modules = getAllModules(), coordinates = {i,j}}
    end
  end
  
  return grid
end

function visalizeGrid(grid)
  for i = 1,#grid do
    for j = 1,#grid[1] do
      local cell = grid[i][j]
      
      print(#cell.modules)
    end
  end
end

function findCommonElements(table1, table2)
    local commonElements = {}

    for _, value1 in ipairs(table1) do
        for _, value2 in ipairs(table2) do
            if value1 == value2 then
                table.insert(commonElements, value1)
                break  -- Break out of the inner loop once a common element is found
            end
        end
    end

    return commonElements
end

function getLowestEntropy(cells)
  local lowestEntropy, bestCell = #getAllModules(), nil
  for _,cell in pairs(cells) do
    local domain = #cells.modules
    if domain < lowestEntropy then
      lowestEntropy, bestCell = domain, cell
    end
  end
  
  return bestCell
end

function collapseCells(grid, cells, modules)
  for _,cell in pairs(cells) do
    local options = findCommonElements(cell.modules, modules)
    
    cell.modules = options
  end
end

function getRandomModule(modules)
  local index = math.random(1, #modules)
  
  return modules[index]
end

function waveCollapse(grid)
  local cells = getAvailableCells(grid)
  
  if #cells == 0 then
    print("Collapse completed")
    return
  end
  
  for cell,neighbours in pairs(cells) do
    collapseCells(grid, neighbours, cell.modules)
  end
  
  local neighbours = cells[1]
  
  local bestCell = getLowestEntropy(neighbours)
  
  bestCell.modules = {getRandomModule()}
end

function start(height, width)
  local grid = createGrid(height, width)
  
  grid[5][5].modules = {1}
  
  repeat
    waveCollapse(grid)
  until #getAvailableCells(grid) == 0
  
  visalizeGrid(grid)
end

start(9,9)