Loop Hero tools

.gitignore

/ocr
/gamedata

ocr.d

import core.stdc.config;
import core.thread.osthread;
import core.time;

import std.algorithm.comparison;
import std.algorithm.iteration;
import std.algorithm.mutation;
import std.algorithm.searching;
import std.algorithm.setops;
import std.algorithm.sorting;
import std.array;
import std.conv;
import std.datetime.systime;
import std.exception;
import std.file;
import std.format.write;
import std.math;
import std.parallelism;
import std.path;
import std.process;
import std.random;
import std.range;
import std.stdio : stderr, File;
import std.string;
import std.typecons;

import deimos.X11.X;
import deimos.X11.Xlib;
import deimos.X11.keysym;

import ae.sys.datamm;
import ae.sys.file;
import ae.sys.sendinput;
import ae.utils.array;
import ae.utils.geometry;
import ae.utils.graphics.color;
import ae.utils.graphics.hls;
import ae.utils.graphics.im_convert;
import ae.utils.graphics.image;
import ae.utils.graphics.view;
import ae.utils.math;
import ae.utils.meta;

pragma(lib, "X11");
pragma(lib, "Xtst");

// ---------------------------------------------------------------------------------------------------------------------

enum gameTitle = "LOOP HERO 1.105 (lin)";

// ---------------------------------------------------------------------------------------------------------------------

Display* dpy;
Window win;
Rect!int geometry;

enum xy_t gameW = 640;
enum xy_t gameH = 360;

Image!ColorIndex getCapture()
{
	static Image!BGRA capture;
	captureWindowRect(win, Rect!int(0, 0, geometry.w, geometry.h), capture);
	static Image!ColorIndex scaled;
	capture
		.warp!(q{x * 2}, q{y * 2})
		.crop(0, 0, capture.w / 2, capture.h / 2)
		.colorMap!(c => BGRA(c.b, c.g, c.r, typeof(c.a).max)) // X11 returns 0 alpha
		.colorMap!(c => colorLookup.get(c, hole))
		.copy(scaled);
	assert(scaled.w == gameW && scaled.h == gameH);
	return scaled;
}

extern (C) int XTestFakeKeyEvent(
    Display*		/* dpy */,
    uint	/* keycode */,
    Bool		/* is_press */,
    c_ulong	/* delay */
);
extern (C) int XTestFakeButtonEvent(
	Display *display,
	uint button,
	Bool is_press, c_ulong delay
);

void moveMouse(int x, int y)
{
	XWarpPointer(dpy, None, win, 0, 0, 0, 0, x, y);
	XSync(dpy, false);
}

void sendButton(int button, bool press)
{
	// XEvent ev2;
	// ev2.type = press ? ButtonPress : ButtonRelease;
	// ev2.xbutton.button = button;
	// ev2.xbutton.same_screen = True;
	// XSendEvent(dpy, win, True/*propagate*/, press ? ButtonPressMask : ButtonReleaseMask, &ev2);
	XTestFakeButtonEvent(dpy, button, press, 1);
	XSync(dpy, false);
}

void sendDoubleClick()
{
	foreach (n; 0 .. 2)
	{
		sendButton(Button1, true);
		Thread.sleep(30.msecs);
		sendButton(Button1, false);
		Thread.sleep(30.msecs);
	}
}

// ---------------------------------------------------------------------------------------------------------------------

alias COLOR = BGRA;

alias ColorIndex = ushort;
COLOR[] colors;
ColorIndex[COLOR] colorLookup;

enum ColorIndex hole = 0;
enum ColorIndex numColors = 564;
static this() { colors ~= COLOR.init; }

ColorIndex internColor(COLOR c)
{
	if (!c.a) return hole;
	if (auto p = c in colorLookup)
		return *p;

	auto index = colors.length.to!ColorIndex;
	colors ~= c;
	colorLookup[c] = index;
	return index;
}

// ---------------------------------------------------------------------------------------------------------------------

struct Sprite
{
	string[] names;
	Image!ColorIndex image;
	size_t numPixels;
}
Sprite[] sprites;
alias SpriteIndex = size_t;

Sprite loadSprite(string fn)
{
	// stderr.writeln("Loading ", fn.baseName);
	scope(failure) stderr.writeln("Error loading ", fn, ":");

	Sprite sprite;
	sprite.names = [fn.baseName.stripExtension];

	auto bmpFn = fn.setExtension(".bmp");
	if (!bmpFn.exists)
	{
		auto png = fn.mapFile(MmMode.read);
		png.contents.parseViaIMConvert!COLOR.toBMP.toFile(bmpFn);
	}
	auto bmp = bmpFn.mapFile(MmMode.read);
	sprite.image = bmp.contents.viewBMP!COLOR.colorMap!internColor.copy();
	sprite.numPixels = sprite.image.pixels.count!(c => c[0] != hole);

	return sprite;
}

void loadSprites()
{
	auto spriteFiles = dirEntries("gamedata/Export_Sprites", "*.png", SpanMode.breadth)
		.map!(de => de.name)
		.array
		.sort
		.map!loadSprite()
		.array;

	stderr.writeln("Deduplicating sprites...");
	spriteFiles.sort!((ref a, ref b) => a.image.pixels < b.image.pixels, SwapStrategy.stable);
	foreach (ref sprite; spriteFiles)
		if (!sprites.length || sprites[$-1].image != sprite.image)
			sprites ~= sprite;
		else
			sprites[$-1].names ~= sprite.names;

	sprites.sort!((a, b) => a.numPixels > b.numPixels, SwapStrategy.stable); // Biggest first
	stderr.writeln("Loaded ", sprites.length, " sprites with ", colors.length, " colors.");
	// enforce(colors.length == numColors, "Color count mismatch");
}

// ---------------------------------------------------------------------------------------------------------------------

enum indexFileName = "index.bin";
enum void* indexAddr = cast(void*)0x1000_0000_0000;
enum size_t indexSize = 0x1000_0000_0000;
void* indexEnd = indexAddr;

void openIndex(File f, bool write)
{
	import core.sys.posix.sys.mman : mmap, PROT_READ, PROT_WRITE;
	import core.sys.linux.sys.mman : MAP_NORESERVE, MAP_SHARED;
	enum MAP_FIXED_NOREPLACE = 0x100000;

	auto result = mmap(
		indexAddr,
		indexSize,
		write ? PROT_READ | PROT_WRITE : PROT_READ,
		MAP_SHARED | MAP_FIXED_NOREPLACE/*|MAP_NORESERVE*/,
		f.fileno, 0);
	errnoEnforce(result == indexAddr);
}
void closeIndex()
{
	import core.sys.posix.sys.mman : munmap;
	munmap(indexAddr, indexSize).I!(ret => ret == 0).errnoEnforce();
}

void* indexAlloc(size_t size)
{
	assert(size % size_t.sizeof == 0);
	auto result = indexEnd;
	indexEnd += size;
	return result;
}
T* indexAlloc(T)() { return cast(T*)indexAlloc(T.sizeof); }
T[] indexAlloc(T)(size_t length) { return (cast(T*)indexAlloc(T.sizeof * length))[0 .. length]; }
T[] indexDup(T)(T[] source) { auto arr = indexAlloc!T(source.length); arr[] = source; return arr; }

/// Matches upon reaching this IndexNode.
struct Match
{
	SpriteIndex spriteIndex;

	/// Delta from sprite x0,y0 to origin (examined pixel)
	/// x0,y0 of sprite + Match.x,y == origin
	xy_t x, y;

	void toString(void delegate(scope const char[]) sink) const
	{
		sink.formattedWrite!"%s @ (%d, %d)"(sprites[spriteIndex].names[0], x, y);
	}

	int opCmp(const ref Match m) const
	{
		if (spriteIndex != m.spriteIndex) return spriteIndex < m.spriteIndex ? 1 : -1;
		if (x != m.x) return x < m.x ? 1 : -1;
		if (y != m.y) return y < m.y ? 1 : -1;
		return 0;
	}
}

struct IndexNode
{
	/// Final matches (no other pixels to check).
	/// If set, the other fields will be unset (leaf node).
	Match[] finalMatches;

	/// Next, check this pixel (delta from origin). Might be negative.
	/// origin + IndexNode.x,y == sample (pixel which will be examined)
	xy_t x, y;

	/// Branches, depending on the pixel's color
	struct Child
	{
		ColorIndex color;
		IndexNode* node; /// Can be null (same as IndexNode with no children or finalMatches)
	}
	Child[] children;
	IndexNode* otherColorNode; /// Go here if the color did not match any children
}
struct Index
{
	IndexNode* root;
	size_t numIndexNodes = 1;
	size_t maxIndexDepth = 0;
	IndexNodeCacheNode nodeCacheRoot;
}
Index* index;

struct IndexNodeCacheNode
{
	IndexNode* indexNode;

	struct Item
	{
		Item* nextSibling;
		Match match;
		IndexNodeCacheNode matchNext;
	}
	Item* items;
}

IndexNode** indexNodeCache(Match[] matches)
{
	auto n = &index.nodeCacheRoot;
	foreach_reverse (ref match; matches)
	{
		IndexNodeCacheNode.Item** i;
		for (i = &n.items; *i; i = &(*i).nextSibling)
			if ((*i).match == match)
				break;
		if (!*i)
		{
			*i = indexAlloc!(IndexNodeCacheNode.Item);
			(*i).match = match;
		}
		assert((*i).match == match);
		n = &(*i).matchNext;
	}
	return &n.indexNode;
}

immutable COLOR[] firstPixelColorsToIgnore = [
	// COLOR(0, 0, 0, 0xFF), // black
];

void makeIndex()
{
	auto f = File(indexFileName ~ ".tmp", "w+b");
	f.truncate(indexSize);
	openIndex(f, true);
	scope(success)
	{
		closeIndex();
		f.close();
		rename(indexFileName ~ ".tmp", indexFileName);
	}

	index = indexAlloc!Index;

	stderr.writeln("Collecting all initial matches...");
	Match[] initialMatches;
	foreach (spriteIndex, ref sprite; sprites)
		foreach (y0; 0 .. sprite.image.h)
			foreach (x0; 0 .. sprite.image.w)
			{
				auto c = sprite.image[x0, y0];
				if (c != hole && !firstPixelColorsToIgnore.canFind(colors[c]))
					initialMatches ~= Match(spriteIndex, x0, y0);
			}

	xy_t[2][] path;

	void expand(size_t depth, ref IndexNode* n, Match[] matches, double progressStart, double progressSize)
	{
		matches.sort();

		stderr.writefln("[%8.4f%%] Expanding at %s: %s", progressStart * 100, depth, matches.length);
		index.maxIndexDepth = max(index.maxIndexDepth, depth);
		if (matches.length <= 1)
		{
			n = indexAlloc!IndexNode();
			n.finalMatches = matches.indexDup;
			return; // Leave matches as final matches
		}

		IndexNode** cache = indexNodeCache(matches);
		if (*cache) { n = *cache; return; }
		scope(success) *cache = n;

		auto maxW = matches.map!((ref match) => sprites[match.spriteIndex].image.w).reduce!max;
		auto maxH = matches.map!((ref match) => sprites[match.spriteIndex].image.h).reduce!max;

		// Find the best next pixel
		xy_t[2] bestXY; size_t bestScore;
	coordSearchOuter:
		foreach (dist; 0 .. maxH + maxW + 1)
			foreach_reverse (y; -maxH + 1 .. maxH)
			  coordSearchInner:
				foreach_reverse (x; -maxW + 1 .. maxW)
					if (abs(x) + abs(y) == dist)
					{
						// Hard-coded optimization: at depth 0, the best coordinate is going to be 0,0.
						if (depth == 0 && (x != 0 || y != 0))
							continue;

						xy_t[2] xy = [x, y];

						// This check prevents repeated attempts to disprove matches
						// by looking at the same pixel several times in one path
						if (path[0 .. depth].canFind(xy))
							continue; // Already looked here

						size_t[numColors] colorCounts; // including `hole`, to indicate out-of-bounds

						foreach (ref match; matches)
						{
							// We are [calculating the viability of] checking: origin + (x,y)
							// We need the sprite coordinates for the match.
							auto sx = match.x + x;
							auto sy = match.y + y;
							auto sprite = &sprites[match.spriteIndex];

							ColorIndex c;
							if (sx < 0 || sy < 0 || sx >= sprite.image.w || sy >= sprite.image.h)
								c = hole;
							else
								c = sprite.image[sx, sy];

							colorCounts[c]++;

							if (c == hole) // Stop early
							{
								// score cannot be possibly bigger than this
								auto scoreUpperBound = (2 + matches.length - colorCounts[hole]) * 10_000_000_000L;
								if (scoreUpperBound < bestScore)
									continue coordSearchInner;
							}
						}

						auto numUsedColors = colorCounts[].count!(count => count > 0);
						if (numUsedColors <= 1)
							continue; // Dead-end - all candidates are in one group

						// 1. Minimize number of holes (as we will carry those results into every branch)
						// 2. Minimize maximum number of matches per color (this is a heuristic)
						// 3. Maximize number of colors
						auto maxMatchesPerColor = colorCounts[hole + 1 .. $].reduce!((a, b) => a > b ? a : b);
						size_t score =
							(1 + matches.length - colorCounts[hole]) * 10_000_000_000L +
							(matches.length - maxMatchesPerColor) * 1_000 +
							numUsedColors;
						// stderr.writeln("  Score at ", x, ",", y, " is ", score, " points");
						if (score > bestScore)
						{
							bestScore = score;
							bestXY = xy;

							if (colorCounts[hole] == 0 && maxMatchesPerColor == 1)
								break coordSearchOuter; // Optimal score found
						}
					}
		stderr.writeln("  Best next coordinate is ", bestXY, " (", bestScore, " points)");

		if (bestScore == 0)
			return; // Dead-end; can't reduce the current match group further, index-user should try from somewhere else

		path.putExpand(depth, bestXY);

		n = indexAlloc!IndexNode();
		n.x = bestXY[0];
		n.y = bestXY[1];

		struct Child
		{
			ColorIndex color;
			Match[] matches;
		}
		Child[] children;
		Match[] allColorMatches;

		foreach (ref match; matches)
		{
			auto sx = match.x + n.x;
			auto sy = match.y + n.y;
			auto sprite = &sprites[match.spriteIndex];

			ColorIndex c;
			if (sx < 0 || sy < 0 || sx >= sprite.image.w || sy >= sprite.image.h)
				c = hole;
			else
				c = sprite.image[sx, sy];

			if (c == hole)
			{
				allColorMatches ~= match;
				continue;
			}

			// if (children.length == 0) stderr.writeln("First match: ", x, ",", y, " ", c);
			auto childIndex = children.countUntil!((ref child) => child.color == c);
			if (childIndex < 0)
			{
				childIndex = children.length;
				children ~= Child(c);
				index.numIndexNodes++;
			}

			children[childIndex].matches ~= match;
		}

		// For faster lookup, put the bigger groups first
		// (they're already likely to be)
		children.sort!((a, b) => a.matches.length > b.matches.length, SwapStrategy.stable);

		n.children = indexAlloc!(IndexNode.Child)(children.length);
		foreach (childIndex, ref child; children)
			n.children[childIndex].color = child.color;

		auto childProgressSize = progressSize / (1 + children.length);

		foreach (childIndex, ref child; children)
			if (child.matches.length + allColorMatches.length == matches.length)
			{} // Dead-end - nothing has been disproved, we are looping in place
			else
				expand(depth + 1, n.children[childIndex].node, child.matches ~ allColorMatches,
					progressStart + (1 + childIndex) * childProgressSize, childProgressSize);

		if (allColorMatches.length)
		{
			if (allColorMatches.length == matches.length)
			{} // Dead-end - nothing has been disproved, we are looping in place
			else
				expand(depth + 1, n.otherColorNode, allColorMatches, progressStart, childProgressSize);
		}
	}

	expand(0, index.root, initialMatches, 0.0, 1.0);
	stderr.writefln("Done expanding: %d nodes, %d maximum depth", index.numIndexNodes, index.maxIndexDepth);
}

void dumpIndex()
{
	auto f = File("index.txt", "wb");
	void dump(IndexNode* n, size_t depth, string name)
	{
		auto prefix = "  ".replicate(depth * 2);
		f.write(prefix, "- ", name, ": ");
		if (!n) { f.writeln("(null)"); return; }
		if (n.finalMatches) f.write(n.finalMatches.length, " final (", n.finalMatches, ")");
		f.writeln();
		if (n.children.length || n.otherColorNode)
		{
			f.writeln(prefix, "  ", n.x, ",", n.y, ":");
			void dumpChild(IndexNode* cn, string name)
			{
				if (cn && cn < n)
					f.writeln(prefix, "    - ", name, " - backreference");
				else
					dump(cn, depth + 1, name);
			}
			foreach (ref child; n.children)
				dumpChild(child.node, "#" ~ colors[child.color].toHex);
			if (n.otherColorNode)
				dumpChild(n.otherColorNode, "[other]");
		}
	}
	dump(index.root, 0, "root");
}

// ---------------------------------------------------------------------------------------------------------------------

struct FoundSprite
{
	size_t spriteIndex;
	xy_t x, y;
	size_t matchedPixels;
}

FoundSprite[] ocr(Image!ColorIndex screen)
{
	FoundSprite[] result;
	bool found;
	do
	{
		found = false;
		FoundSprite bestFinding;

		foreach (oy; 0 .. screen.h)
		  coordLoop:
			foreach (ox; 0 .. screen.w)
			{
				auto n = index.root;

				while (true)
				{
					if (n.children.length == 0)
						break;

					auto x = ox + n.x;
					auto y = oy + n.y;
					ColorIndex c;
					if (x < 0 || y < 0 || x >= screen.w || y >= screen.h)
						c = hole;
					else
						c = screen[x, y];

					if (c == hole)
						continue coordLoop; // Can't proceed due to hole - will try from another origin instead

					auto childIndex = n.children.countUntil!((ref child) => child.color == c);
					if (childIndex < 0)
						n = n.otherColorNode;
					else
						n = n.children[childIndex].node;
					if (!n)
						continue coordLoop; // Dead end
				}

			checkMatch:
				foreach (ref match; n.finalMatches)
				{
					auto sprite = &sprites[match.spriteIndex];
					auto sx0 = ox - match.x;
					auto sy0 = oy - match.y;

					size_t matchedPixels;
					foreach (sy; 0 .. sprite.image.h)
						foreach (sx; 0 .. sprite.image.w)
						{
							auto cSpr = sprite.image[sx, sy];
							if (cSpr == hole) continue;

							auto cx = sx0 + sx;
							auto cy = sy0 + sy;
							if (cx < 0 || cy < 0 || cx >= screen.w || cy >= screen.h)
								continue checkMatch;

							auto cScr = screen[cx, cy];
							if (cScr == hole) continue;

							if (cSpr != cScr)
								continue checkMatch;

							matchedPixels++;
						}

					if (matchedPixels > bestFinding.matchedPixels)
					{
						bestFinding = FoundSprite(match.spriteIndex, sx0, sy0, matchedPixels);
						stderr.writeln(
							"New best finding: ", sprite.names,
							" at ", bestFinding.x, ",", bestFinding.y,
							" (", bestFinding.matchedPixels, "/", sprite.numPixels, ")");
					}
				}
		}

		if (bestFinding.matchedPixels > 0)
		{
			auto sprite = &sprites[bestFinding.spriteIndex];
			stderr.writeln(
				"Found ", sprite.names,
				" at ", bestFinding.x, ",", bestFinding.y,
				" (", bestFinding.matchedPixels, "/", sprite.numPixels, ")");

			found = true;
			result ~= bestFinding;

			// Erase and forget
			foreach (sy; 0 .. sprite.image.h)
				if (bestFinding.y + sy < screen.h)
					foreach (sx; 0 .. sprite.image.w)
						if (bestFinding.x + sx < screen.w)
						{
							if (sprite.image[sx, sy] != hole)
								screen[bestFinding.x + sx, bestFinding.y + sy] = hole;
						}
		}
	}
	while (found);

	return result;
}

// ---------------------------------------------------------------------------------------------------------------------

void play()
{
	loadSprites();
	if (!indexFileName.exists)
		makeIndex();

	auto f = File(indexFileName, "rb");
	openIndex(f, false);
	index = cast(Index*)indexAddr;

	// dumpIndex();

	dpy = getDisplay();
	win = findWindowByName(gameTitle);
	geometry = getWindowGeometry(win);

	auto c = getCapture();
	ocr(c);
	c.colorMap!(c => colors[c]).colorMap!(c => RGBA(c.r, c.g, c.b, c.a)).toPNG.toFile("ocr.png");
}

import ae.utils.funopt;
import ae.utils.main;
mixin main!(funopt!play);

// "