poke(0x5F2D, 3)

--
-- constants
--

local BLUE = 1
local YELLOW = 2

--
-- Utils
--

function idx_xy(idx, width)
	return (idx - 1) % width + 1, (idx - 1) \ width + 1
end

function xy_idx(x, y, width)
	return ((y - 1) * width) + x
end

function map(tbl, f)
	local t = {}
	for k, v in pairs(tbl) do
		t[k] = f(v)
	end
	return t
end

function filter(tbl, f, keepindex)
	local ret = {}
	for k, v in pairs(tbl) do
		if f(v) then
			if keepindex
				then ret[k] = v
				else add(ret, v)
			end
		end
	end
	return ret
end

function slice(tbl, first, last, step)
	local sliced = {}
	for i = (first or 1), (last or #tbl), (step or 1) do
		sliced[(#sliced + 1)] = tbl[i]
	end
	return sliced
end

local _rectfill = rectfill
function rectfill(x, y, w, h, col)
	_rectfill(x, y, x+w, y+h, col)
end

local _count = count
function count(tbl, p)
	if type(p) != "function" then return _count(tbl, p) end
	local c = 0
	for v in all(tbl) do
		if p(v) then c+=1 end
	end
	return c
end

function find(tbl, o)
	local indices = {}
	for k,v in ipairs(tbl) do
		if v == o then add(indices, k) end
	end
	return indices
end

-- shallow copy
function copy(tbl)
	local copy = {}
	for v in all(tbl) do add(copy, v) end
	return copy
end

-- shallow equality
function equal(tbl1, tbl2)
	for k, _ in ipairs(tbl1) do
		if tbl1[k] ~= tbl2[k] then return false end
	end
	return true
end

function tostring(any)
  if (type(any)~="table") return tostr(any)
  local str = "{"
  for k,v in pairs(any) do
    if (str~="{") str=str..","
    str=str..tostring(k).."="..tostring(v)
  end
  return str.."}"
end

--
-- Board
--

local Board = {}
function Board.new()
	local debug = false
	local width = 10
	local t = [[
0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 2, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2, 2, 0, 1, 0, 2, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 2, 0, 0
	]]
	a = {foo = 3}
	local tiles = split(t)

	return {
		reset = function(self)
			local t = [[
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0,
				0,0,0,0,0,0,0,0,0,0
			]]
			tiles = split(t)
		end,

		getWidth = function()
			return width
		end,

		getTile = function(self, idx)
			return tiles[idx]
		end,

		-- returns a COPY of the tiles array
		getTilesCopy = function(self)
			return copy(tiles)
		end,

		-- overwrites the whole tiles array
		setTiles = function(self, newtiles)
			assert(#newtiles == #tiles, "New tiles array must have a length of " .. #tiles)
			tiles = copy(newtiles)
		end,

		idx_xy = function(self, idx)
			return idx_xy(idx, width)
		end,

		xy_idx = function(self, x, y)
			return xy_idx(x, y, width)
		end,

		fill = function(self, idx, color, invert)
			if invert then
				color = color == YELLOW and BLUE or YELLOW
			end
			tiles[idx] = color
		end,

		getRows = function(self)
			local ret = {}
			for i = 1, width do
				add(ret, slice(tiles, ((i - 1) * width) + 1, i * width))
			end
			return ret
		end,

		getCols = function(self)
			local ret = {}
			local rows = self.getRows(self)
			for i = 1, width do
				add(ret, map(rows, function(v) return v[i] end))
			end
			return ret
		end,

		isComplete = function(self)
			return count(tiles, 0) == 0
		end,

		isValid = function(self)
			local rows = self:getRows()
			for y,row in ipairs(rows) do
				-- check count
				if count(row, BLUE) ~= count(row, YELLOW) then
					if (debug) printh("uneven count on row "..y)
					return false
				end
				-- check identical lines
				for k,other in ipairs(rows) do
					if equal(other, row) and other ~= row then
						if (debug) printh("equal rows "..k)
						return false
					end
				end
				-- check triples
				if self:countConsecutives(row) > 2 then
					if (debug) printh("triples")
					return false
				end
			end

			local cols = self:getCols()
			for col in all(cols) do
				-- check count
				if count(col, BLUE) ~= count(col, YELLOW) then
					if (debug) printh("uneven count")
					return false
				end
				-- check identical lines
				for other in all(cols) do
					if equal(other, col) and other ~= col then
						if (debug) printh("equal cols")
						return false
					end
				end
				-- check triples
				if self:countConsecutives(col) > 2 then
					if (debug) printh("triples")
					return false
				end
			end
			return true
		end,

		countConsecutives = function(self, line)
			local count = 0
			local last = 0
			for v in all(line) do
				if v ~= last then
					last = 0
				else
					last = v
					count += 1
				end
			end
			return count
		end,

		-- Returns the index of a random zero tile
		getRandomZero = function(self)
			assert(count(tiles, 0) > 0, "No zero left")
			local zeroes = filter(tiles, function(v) return v == 0 end, true)
			local z = {}
			for k,v in pairs(zeroes) do
				add(z, k)
			end
			return rnd(z)
		end,

		-- Returns the index of a random non-zero tile
		getRandomNonZero = function(self)
			assert(count(tiles, 0) < #tiles, "All zeroes")
			local numbers = filter(tiles, function(v) return v ~= 0 end, true)
			local z = {}
			for k,v in pairs(numbers) do
				add(z, k)
			end
			return rnd(z)
		end,

		tostring = function(self)
			local str = ''
			for v in all(tiles) do
				str ..= ", " .. v
			end
			return str
		end,

		-- Solves 1 step of the board
		-- Returns "valid" if it solved it without guessing
		-- Returns "invalid" if the board cannot be solved
		solveStep = function(self, random)
			local zeroes = count(tiles, 0)
			self:surroundDoubles()
			self:splitTriples()
			self:fillLines()
			self:noIdenticalLines()
			local changed = zeroes ~= count(tiles, 0)
			if not changed and random and not self:isComplete() then
				-- Set a random color
				local z = self:getRandomZero()
				self:fill(z, rnd({BLUE, YELLOW}))
				if (debug) printh("!!!!!!!!!!!!!!!!! RANDOM FILL AT " .. z)
				return "invalid"
			end
			return (changed or self:isComplete()) and "valid" or "invalid"
		end,

		surroundDoubles = function(self)
			for idx,v in ipairs(tiles) do
				local x,y = self:idx_xy(idx)
				if v == 0 then
					local neighbors = {}
					-- 2 tiles on the left
					if x >= 3 then
						add(neighbors, {idx, idx-1, idx-2})
					end
					-- 2 tiles on the right
					if x <= width-2 then
						add(neighbors, {idx, idx+1, idx+2})
					end
					-- 2 tiles on top
					if y >= 3 then
						add(neighbors, {idx, idx-width, idx - width*2})
					end
					-- 2 tiles under
					if y <= width-2 then
						add(neighbors, {idx, idx+width, idx + width*2})
					end

					-- only keep pairs that are identical (and not 0)
					neighbors = filter(neighbors, function (o) return tiles[o[2]] == tiles[o[3]] and tiles[o[2]] ~= 0 end)

					-- do the surrounding
					for item in all(neighbors) do
						if item[1] then
							if (debug) printh("Surrounding at " .. item[1])
							self:fill(item[1], tiles[item[2]], true)
						end
					end
				end
			end
		end,

		splitTriples = function(self)
			for idx, col in ipairs(tiles) do
				local x,y = self:idx_xy(idx)
				if col == 0 then

					if x > 1 and x < width then
						-- check horizontal
						local prev = tiles[idx-1]
						local next = tiles[idx+1]
						if prev ~= 0 and prev == next then
							if (debug) printh("Splitting at " .. idx)
							self:fill(idx, prev, true)
						end
					end

					if y>1 and y < width then
						-- check vertical
						local prev = tiles[idx-width]
						local next = tiles[idx+width]
						if prev ~= 0 and prev == next then
							if (debug) printh("Splitting at " .. idx)
							self:fill(idx, prev, true)
						end
					end
				end
			end
		end,

		fillLines = function(self)
			local rows = self:getRows()
			local cols = self:getCols()

			-- rows
			for y,row in ipairs(rows) do
				local a = count(row, BLUE)
				local b = count(row, YELLOW)
				if a ~= b then
					if a == width/2 then self:fillRow(y, YELLOW) end
					if b == width/2 then self:fillRow(y, BLUE) end
				end
			end

			-- columns
			for x,col in ipairs(cols) do
				local a = count(col, BLUE)
				local b = count(col, YELLOW)
				if a ~= b then
					if a == width/2 then self:fillCol(x, YELLOW) end
					if b == width/2 then self:fillCol(x, BLUE) end
				end
			end
		end,

		fillRow = function(self, y, color)
			if (debug) printh("Filling line " .. y .. " in " .. (color == BLUE and "blue" or "yellow"))
			local idx = self:xy_idx(1, y)
			for i = idx, (idx+width-1) do
				if self:getTile(i) == 0 then
					self:fill(i, color)
				end
			end
		end,

		fillCol = function(self, x, color)
			if (debug) printh("Filling column " .. x .. " in " .. (color == BLUE and "blue" or "yellow"))
			local idx = self:xy_idx(x, 1)
			for i = idx, #tiles, width do
				if self:getTile(i) == 0 then
					self:fill(i, color)
				end
			end
		end,

		noIdenticalLines = function(self)
			-- columns
			local cols = self:getCols()
			for x,col in ipairs(cols) do
				if count(col, 0) == 2 and count(col, BLUE) == count(col, YELLOW) then
			    local y1, y2 = unpack(find(col, 0))
					local ab = copy(col) ab[y1] = BLUE ab[y2] = YELLOW
					local ba = copy(col) ba[y1] = YELLOW ba[y2] = BLUE
					-- Check if a dupe exists
					for x2,col in ipairs(cols) do
						if equal(col, ab) then
							if debug then
								printh("No-dupe at col " .. x .. " from col " .. x2)
								printh("  " .. self:xy_idx(x,y1) .. " in yellow")
								printh("  " .. self:xy_idx(x,y2) .. " in blue")
							end
							self:fill(self:xy_idx(x,y1), YELLOW)
							self:fill(self:xy_idx(x,y2), BLUE)
							goto continue
						elseif equal(col, ba) then
							if debug then
								printh("No-dupe at col " .. x .. " from col " .. x2)
								printh("  " .. self:xy_idx(x,y1) .. " in blue")
								printh("  " .. self:xy_idx(x,y2) .. " in yellow")
							end
							self:fill(self:xy_idx(x,y1), BLUE)
							self:fill(self:xy_idx(x,y2), YELLOW)
							goto continue
						end
					end
				end
				::continue::
			end

			-- rows
			local rows = self:getRows()
			for y,row in ipairs(rows) do
				if count(row, 0) == 2 and count(row, BLUE) == count(row, YELLOW) then
			    local x1, x2 = unpack(find(row, 0))
					local ab = copy(row) ab[x1] = BLUE ab[x2] = YELLOW
					local ba = copy(row) ba[x1] = YELLOW ba[x2] = BLUE
					-- Check if a dupe exists
					for y2,row in ipairs(rows) do
						if equal(row, ab) then
							if debug then
								printh("No-dupe at row " .. y .. " from row " .. y2)
								printh("  " .. self:xy_idx(x1,y) .. " in yellow")
								printh("  " .. self:xy_idx(x2,y) .. " in blue")
							end
							self:fill(self:xy_idx(x1,y), YELLOW)
							self:fill(self:xy_idx(x2,y), BLUE)
							goto continue
						elseif equal(row, ba) then
							if debug then
								printh("No-dupe at row " .. y .. " from row " .. y2)
								printh("  " .. self:xy_idx(x1,y) .. " in blue")
								printh("  " .. self:xy_idx(x2,y) .. " in yellow")
							end
							self:fill(self:xy_idx(x1,y), BLUE)
							self:fill(self:xy_idx(x2,y), YELLOW)
							goto continue
						end
					end
				end
				::continue::
			end
		end,

		draw = function(self)
			local w = 10
			local padding = 1
			local margin = 4
			for k,v in ipairs(tiles) do
				local x,y = self:idx_xy(k)
				local color = v == BLUE and 9 or v == YELLOW and 3 or 1
				rectfill(
					margin + (x-1)*w + (x-1)*(padding+1),
					margin + (y-1)*w + (y-1)*(padding+1),
					w, w, color)
					if debug then
						print(k,
							margin + (x-1)*w + (x-1)*(padding+1),
							margin + (y-1)*w + (y-1)*(padding+1), 8)
					end
			end
		end
	}
end

--
-- main loop
--

local board = Board.new()
local create = true

function _init()
	cls()
	printh(" ")
	printh(" ")
	-- srand(42)

	if create then
		-- Create a board
		repeat
			board:reset()
			repeat
				board:solveStep(true)
			until board:isComplete()
		until board:isValid()

		-- Remove tiles until randomness is unavoidable
		local previous = {board:getTilesCopy()} -- initial state
		local invalidcount = 0
		while true do
			-- remove a random tile
			board:setTiles(previous[#previous])
			local i = board:getRandomNonZero()
			board:fill(i, 0)
			add(previous, board:getTilesCopy())

			-- try to solve the board
			local solved = ""
			repeat
				solved = board:solveStep()
			until board:isComplete() or solved == "invalid"
			if board:isComplete() then
				invalidcount = 0
			end
			if solved == "invalid" then
				deli(previous)
				invalidcount += 1
			end
			if invalidcount == 100 then
				break
			end
		end -- end while

		board:setTiles(previous[#previous])
		printh(board:tostring())
	end
end

function _update60()
	if not create then
		board:solveStep()
	end
end

function _draw()
	cls()
	board:draw()
	circfill(stat(32), stat(33), 0, 7)
	circ(stat(32), stat(33), 1, 0)
end