1 files changed, 90 insertions, 1 deletions

diff --git a/img/shitty_disp.lua b/img/shitty_disp.lua
index f5ebfb9..3f060d9 100644
--- a/img/shitty_disp.lua
+++ b/img/shitty_disp.lua
@@ -6,7 +6,7 @@ local function color_dist(c1,c2)
 	return total
 end
 
-function closest_palette_color(R,c)
+local function closest_palette_color(R,c)
 	local closest = nil
 	local closest_dist = 99999999
 	for i=0,15 do
@@ -20,6 +20,93 @@ function closest_palette_color(R,c)
 	return closest
 end
 
+local function average_color(pxlist)
+	local count = #pxlist
+	local tr,tg,tb = 0,0,0
+	for _,px in ipairs(pxlist) do
+		tr = tr + px[1]
+		tg = tg + px[2]
+		tb = tb + px[3]
+	end
+	return {tr/count,tg/count,tb/count}
+end
+
+local function sortby_r(pxlist) table.sort(pxlist,function(a,b) return a[1]<b[1] end) end
+local function sortby_g(pxlist) table.sort(pxlist,function(a,b) return a[2]<b[2] end) end
+local function sortby_b(pxlist) table.sort(pxlist,function(a,b) return a[3]<b[3] end) end
+
+local function median_cut(pxlist,remaining)
+	remaining = remaining or 4
+	if remaining <= 0 then
+		return {average_color(pxlist)}
+	end
+	local maxr,maxg,maxb = -1,-1,-1
+	local minr,ming,minb = 999,999,999
+	local max,min = math.max,math.min
+	for _,px in ipairs(pxlist) do
+		maxr = max(maxr,px[1])
+		maxg = max(maxg,px[2])
+		maxb = max(maxb,px[3])
+		minr = min(minr,px[1])
+		ming = min(ming,px[2])
+		minb = min(minb,px[3])
+	end
+	local ranger,rangeg,rangeb = maxr-minr,maxg-ming,maxb-minb
+	if ranger > rangeg then
+		if ranger > rangeb then
+			-- red
+			sortby_r(pxlist)
+		else
+			-- blue
+			sortby_b(pxlist)
+		end
+	else
+		if rangeg > rangeb then
+			-- green
+			sortby_g(pxlist)
+		else
+			-- blue
+			sortby_b(pxlist)
+		end
+	end
+	local count = #pxlist
+	local hcount = math.floor(count/2)
+
+	local lhalf,uhalf = {},{}
+	table.move(pxlist, 1,hcount, 1, lhalf)
+	table.move(pxlist, hcount+1,count, 1, uhalf)
+
+	local lcols = median_cut(lhalf,remaining-1)
+	local ucols = median_cut(uhalf,remaining-1)
+	
+	return table.move(ucols,1,#ucols,#lcols+1,lcols)
+end
+
+local function reset_palette(R)
+	for i=0,15 do
+		local cx = 2^i
+		R.setPaletteColor(cx,R.nativePaletteColor(cx))
+	end
+end
+
+local function improve_palette(R,img,w,h,imgx,imgy)
+	local pxlist = {}
+	local idx=1
+	for px=imgx,imgx+w-1 do
+		for py=imgy,imgy+h-1 do
+			pxlist[idx] = img[py][px]
+			idx = idx + 1
+		end
+	end
+	local palette = median_cut(pxlist)
+	assert(#palette==16,#palette.." doesn't equal 16")
+	for i=0,15 do
+		local pc = palette[i+1]
+		R.setPaletteColor(2^i,table.unpack(pc))
+	end
+end
+	
+
 
 local function disp(R,img,scrx,scry,w,h,imgx,imgy)
 	scrx = scrx or 1
@@ -29,6 +116,8 @@ local function disp(R,img,scrx,scry,w,h,imgx,imgy)
 	imgx = imgx or 1
 	imgy = imgy or 1
 
+	improve_palette(R,img,w,h,imgx,imgy)
+
 	for offx=0,w-1 do
 		for offy=0,h-1 do
 			R.setCursorPos(scrx+offx,scry+offy)