#include "common/pch.h"
#include "ThsCodeOCR.h"
#include "win32_screen_shots.h"
#include "ImgDivider.h"

//
using namespace tesseract;

cv::Mat ThsCodeOCR::grayImage(cv::Mat src) {
	cv::Mat grayImage;
	if (src.channels() == 3)
	{
		cvtColor(src, grayImage, cv::COLOR_RGB2GRAY);
	}
	else if (src.channels() == 4) {
		cvtColor(src, grayImage, cv::COLOR_RGBA2GRAY);
	}
	else if (src.channels() == 1) {
		return src;
	}
	return grayImage;
}

ThsCodeOCR::ThsCodeOCR()
{
	tessApi = new tesseract::TessBaseAPI();
	// ÉèÖÃÓïÑÔ
	if (tessApi->Init("tessdata", "eng")) {
		std::cerr << "Could not initialize tesseract.\n";
		exit(1);
	}


}

ThsCodeOCR::~ThsCodeOCR()
{
	if (tessApi != nullptr) {
		// ÇåÀí
		tessApi->End();
		delete tessApi;
	}
}

cv::Mat ThsCodeOCR::splitForCodeImg(cv::Mat src)
{
	int rows = src.rows;
	int cols = src.cols;
	int full_row = -1;
	int start_row = -1;
	int  end_row = -1;
	for (int r = 0; r < rows; r++) {
		if (ImgDivider::isRowFull(src, r, 0, cols / 2, 1, 14, 17)) {
			full_row = r;
			if (start_row >= 0 && r - start_row > 10) {
				end_row = r - 1;
				break;
			}
		}
		else {
			if (full_row > -1) {
				full_row = -1;
				if (start_row < 0) {
					start_row = r;
				}
			}
		}
	}

	if (end_row < 0)
		end_row = rows - 1;
	if (start_row < 0)
		throw string("ûÕÒµ½ÉÏ·Ö¸îÏß");

	int  start_col = cols - 1;
	for (int c = cols*3/4; c >= 0; c--) {
		if (ImgDivider::isColFull(src, c, (start_row+ end_row)/2, end_row, 2, 14, 17)) {
			start_col = c + 1;
			break;
		}
	}

	int start_index = -1;
	int end_index = -1;
	list<POINT> codes_pos;
	int max_color = 0;
	// ¼ÆËãÓұߺóÈý·ÖÖ®Ò»µÄ×î´óÏñËØ
	for (int r = start_row; r < end_row; r++) {
		for (int c = cols - 1; c > start_col + (cols - start_col) / 3; c--) {
			if (src.ptr<uchar>(r)[c] > max_color)
				max_color = src.ptr<uchar>(r)[c];
		}
	}
	int thresh_hold_color = max_color * 2 / 3;

	for (int col = cols - 1; col > start_col; col--) {
		if (!ImgDivider::isColEmpty(src, col, start_row, end_row, thresh_hold_color)) {
			if (start_index < 0)
			{
				start_index = col;
			}
			end_index = col;
		}
		else {
			if (end_index >= 0 && start_index >= 0) {
				codes_pos.push_back(POINT({ end_index - 1, start_index + 1 }));
				end_index = -1;
				start_index = -1;
				if (codes_pos.size() >= 6) {
					break;
				}
			}
		}
	}

	if (codes_pos.size() < 6) {
		throw string("ûÓÐÕÒµ½6¸öÊý×Ö");
	}

	codes_pos.sort([](POINT p1, POINT p2) {
		return p1.x < p2.x;
		});
	start_col = codes_pos.begin()->x;
	int end_col = codes_pos.back().y;
	return cv::Mat(src, cv::Rect(start_col, start_row, end_col - start_col + 1, end_row - start_row + 1));
}

string ThsCodeOCR::ocrCode(cv::Mat src)
{


	Pix* pix = pixCreate(src.cols, src.rows, 8);
	for (int y = 0; y < src.rows; ++y) {
		for (int x = 0; x < src.cols; ++x) {
			pixSetPixel(pix, x, y, src.ptr<uchar>(y)[x]);
		}
	}

	// ÉèÖÃͼÏñ
	tessApi->SetImage(pix);
	// »ñÈ¡Îı¾
	char* outText = tessApi->GetUTF8Text();
	string text = string(outText);
	std::cout << "OCR output:\n" << outText << "\n";
	pixDestroy(&pix);
	delete[] outText;
	text = StringUtil::trim(text);
	// È¡×îºó6λ
	if (text.length() > 6) {
	  text = text.substr(text.length()-6, 6);
	}
	return text;
}

string ThsCodeOCR::ocrCode()
{
	RECT rect;
	cv::Mat mat = ThsCodeOCR::capture(ConfigUtil::getOCRWindowParams());
	mat = ThsCodeOCR::splitForCodeImg(mat);
	string code = ocrCode(mat);
	return code;
}


HWND ThsCodeOCR::getTHSWindow()
{
	list<HWND> hwnds = Win32Util::searchWindow("ͬ»¨Ë³");
	for (list<HWND>::iterator e = hwnds.begin(); e != hwnds.end(); ++e) {
		HWND hwnd = *e;
		HWND child = NULL;
		for (int i = 0; i < 20; i++) {

			child = FindWindowExA(hwnd, child, NULL, NULL);
			if (child > 0) {
				RECT rect;
				GetWindowRect(child, &rect);
				if (rect.right - rect.left > 1000 && rect.bottom - rect.top > 800) {
					return child;
				}
			}
		}
	}
	return HWND();
}

cv::Mat ThsCodeOCR::capture(OCRWindowParams params)
{
	HWND hwnd = getTHSWindow();
	if (hwnd <= 0) {
		throw string("ûÓлñÈ¡µ½Í¬»¨Ë³´°¿Ú");
	}
	RECT rect;
	Win32Util::getWindowRect(hwnd, &rect);
	//[Éϱ߾à, ¿í¶È, ¸ß¶È, Ëõ·Å±¶Êý]
	int owidth = rect.right - rect.left + 1;
	RECT area;
	area.top = params.top;
	area.left = owidth - params.width;
	area.bottom = area.top + params.height - 1;
	area.right = owidth - 2;

	cv::Mat mat = _wss::get_window_bits(hwnd, &area);
	// »Ò¶Èת»»
	mat = grayImage(mat);
	return mat;
}