#include <iostream>
#include <iomanip>
#include <cstdio>
#include <cstdlib>
#include <cassert>
#include <set>
#include <unistd.h>
#include <tesseract/baseapi.h>
#include <leptonica/allheaders.h>
#include "lodepng.h"

using namespace std;

struct __attribute__((packed)) Clr {
	uint8_t r, g, b;

	uint8_t avg() const {
		return ((int)r + g + b) / 3;
	}

	int bright() const {
		return (int)r + g + b;
	}

	float blueness() const {
		return ((int)b + 1) / (((int)r + g) / 2.0f + 1);
	}
};

bool operator<(const Clr &a, const Clr &b) {
	return a.bright() < b.bright();
}

ostream& operator<<(ostream &os, const Clr &c) {
	return os << (unsigned)c.r << " " << (unsigned)c.g << " " << (unsigned)c.b;
}

void writeimg(const Clr *g, int gw, int gh, const string &fname) {
	assert(lodepng_encode24_file(fname.data(), (uint8_t*)g, gw, gh) == 0);
}

tesseract::TessBaseAPI *tessapi = nullptr;

char tesseract_recog(const Clr *g, int gw, int gh) {
	if (!tessapi) {
		tessapi = new tesseract::TessBaseAPI();
		// Initialize tesseract-ocr with English, without specifying tessdata path
		if (tessapi->Init(nullptr, "eng")) {
			fprintf(stderr, "Could not initialize tesseract.\n");
			exit(1);
		}
	}

	const char *fname = ".tesseract_recog.png";
	writeimg(g, gw, gh, fname);

	// Open input image with leptonica library
	Pix *image = pixRead(fname);
	unlink(fname);
	pixSetResolution(image, 70, 70);
	tessapi->SetImage(image);
	// Get OCR result
	char *outText = tessapi->GetUTF8Text();
	char c = outText[0];
	delete[] outText;

	// Destroy used object and release memory
	// tessapi->End();
	pixDestroy(&image);

	return c;
}

void tesseract_finish() {
	tessapi->End();
	delete tessapi;
}

Clr *img;
int w, h;

template <bool (*F)(const Clr &c)>
int flood(int x, int y, bool *mark) {
	if (!F(img[w*y+x])) return 0;
	mark[w*y+x] = true;
	int t = 1;
	if (x > 0 && !mark[w*y+x-1]) t += flood<F>(x-1, y, mark);
	if (y > 0 && !mark[w*(y-1)+x]) t += flood<F>(x, y-1, mark);
	if (x < w-1 && !mark[w*y+x+1]) t += flood<F>(x+1, y, mark);
	if (y < h-1 && !mark[w*(y+1)+x]) t += flood<F>(x, y+1, mark);
	return t;
}

void cmd_histo() {
	int cnt[256] = {};

	for (int y = 0; y < h; y++) {
		int avg = 0;
		for (int x = 0; x < w; x++) avg += img[w*y+x].avg();
		avg /= w;
		cnt[avg]++;
	}
	for(int i = 0; i < 256; i++) {
		cout << setw(3) << i;
		for (int j = 0; j < cnt[i]; j++) cout << "*";
		cout << endl;
	}
}

void cmd_clrlist() {
	set<Clr> s;

	for (int i = 0; i < w * h; i++) {
		s.insert(img[i]);
	}

	for (const Clr &c : s) {
		cout << c << "  " << c.bright() << " " << c.blueness() << endl;
	}
}

bool clr_interest(const Clr &c) {
	return c.bright() <= 300 && c.blueness() <= 1.6;
}

Clr* subimage(const Clr *orig, int ow, int sx, int sy, int sw, int sh) {
	Clr *sub = new Clr[sw * sh];
	for (int y = 0; y < sh; y++) {
		memcpy(sub + sw * y, orig + ow * (y + sy) + sx, 3 * sw);
	}
	return sub;
}

void cmd_recog() {
	for (int i = 0; i < w * h; i++) {
		if (!clr_interest(img[i])) {
			img[i] = {255, 255, 255};
		}
	}

	// writeimg(img, w, h, "out1.png");

	bool *mark = new bool[w*h]();
	bool *mark2 = new bool[w*h];
	int maxn = -1, maxi = -1;

	for (int y = 0; y < h; y++) {
		for (int x = 0; x < w; x++) {
			if (!mark[w*y+x]) {
				int n = flood<clr_interest>(x, y, mark);
				if (n > maxn) {
					maxn = n;
					maxi = w*y + x;
				}
				if (0 < n && n < 300) {
					memset(mark2, 0, w*h*sizeof(bool));
					flood<clr_interest>(x, y, mark2);
					for (int i = 0; i < w*h; i++) {
						if (mark2[i]) img[i] = {255, 255, 255};
					}
				}
			}
		}
	}

	delete[] mark2;

	memset(mark, 0, w*h*sizeof(bool));
	flood<clr_interest>(maxi%w, maxi/w, mark);

	int minx = w, maxx = 0, miny = h, maxy = 0;

	for (int y = 0; y < h; y++) {
		for (int x = 0; x < w; x++) {
			if (mark[w*y+x]) {
				img[w*y+x] = {255, 255, 255};
				if (x < minx) minx = x;
				if (y < miny) miny = y;
				if (x > maxx) maxx = x;
				if (y > maxy) maxy = y;
			} else if (img[w*y+x].bright() != 3*255) {
				img[w*y+x] = {0, 0, 0};
			}
		}
	}

	delete[] mark;

	/* int w2 = maxx - minx + 1;
	int h2 = maxy - miny + 1;

	Clr *img2 = new Clr[w2*h2];
	for (int y = 0; y < h2; y++) {
		for (int x = 0; x < w2; x++) {
			img2[w2*y+x] = img[w*(y+miny)+x+minx];
		}
	}

	writeimg(img2, w2, h2, "out.png");

	delete[] img2; */

	int result[81];
	for (int yi = 0; yi < 9; yi++) {
		for (int xi = 0; xi < 9; xi++) {
			int sx = minx + (maxx - minx) * xi / 9;
			int sy = miny + (maxy - miny) * yi / 9;
			int sw = (maxx - minx) / 9;
			int sh = (maxy - miny) / 9;
			Clr *sub = subimage(img, w, sx, sy, sw, sh);
			// writeimg(sub, sw, sh, "sub" + to_string(yi) + "_" + to_string(xi) + ".png");
			char c = tesseract_recog(sub, sw, sh);
			if (c == '\0') result[9*yi+xi] = -1;
			else result[9*yi+xi] = c - '0';
		}
	}

	for (int yi = 0; yi < 9; yi++) {
		if (yi > 0 && yi % 3 == 0) cout << endl;
		for (int xi = 0; xi < 9; xi++) {
			if (xi > 0 && xi % 3 == 0) cout << ' ';
			if (xi > 0) cout << ' ';
			cout << ".0123456789"[result[9*yi+xi] + 1];
		}
		cout << endl;
	}

	tesseract_finish();
}

int main(int argc, char **argv){
	if (argc != 3) {
		cout << "Usage: " << argv[0] << " <cmd> <file.png>" << endl;
		return 1;
	}

	const char *cmd = argv[1];
	const char *srcfile = argv[2];

	uint8_t *imgdata = nullptr;

	assert(lodepng_decode24_file(&imgdata, (unsigned*)&w, (unsigned*)&h, srcfile) == 0);
	img = (Clr*)imgdata;

	if (strcmp(cmd, "h") == 0) cmd_histo();
	else if (strcmp(cmd, "l") == 0) cmd_clrlist();
	else if (strcmp(cmd, "r") == 0) cmd_recog();
	else cout << "Unknown command" << endl;
}