/* Stephane Charette, stephanecharette@gmail.com
 * License:  CC0, no rights reserved, PUBLIC DOMAIN.
 * Tool to generate stone barcode synthetic images to train a neural net.
 * See:  https://www.ccoderun.ca/programming/2019-08-18_Darknet_training_images/
 */

#include <map>
#include <string>
#include <opencv2/opencv.hpp>
#include <zint.h>


/* Map a simple number to a OpenCV Mat object.  Used to preload all the background images.
 * See load_all_background_images() for details.
 */
static std::map<int, cv::Mat> all_backgrounds;


void load_all_background_images()
{
	if (all_backgrounds.empty())
	{
		/* We should have a number of background images that we can use.  They are all named the same way, starting at 0
		 * and increasing until there are no more images.  Read them all in, and store them in the "all_backgrounds" map
		 * so we can easily retrieve them later on when building the full images.  The background filenames look like this:
		 *
		 * ../backgrounds/background_0.jpg
		 * ../backgrounds/background_1.jpg
		 * ../backgrounds/background_2.jpg
		 * ../backgrounds/background_3.jpg
		 * ...etc...
		 */
		int counter = 0;
		while (true)
		{
			const std::string filename = "../backgrounds/background_" + std::to_string(counter) + ".jpg";
			cv::Mat mat = cv::imread(filename);
			if (mat.empty())
			{
				// assume we have no more backgrounds to read
				break;
			}

			std::cout << "loaded background image " << filename << std::endl;
			all_backgrounds[counter] = mat;
			counter ++;
		}
	}

	std::cout << "number of background images loaded: " << all_backgrounds.size() << std::endl;

	return;
}


cv::Mat get_random_background_image()
{
	cv::Mat background;

	int r = std::rand() % 100;
	if (r < 5 || all_backgrounds.empty())
	{
		// every once in a while use a plain grey background instead of one of our preloaded background images
		const int red	= 224 + std::rand() % 32;
		const int green	= 224 + std::rand() % 32;
		const int blue	= 224 + std::rand() % 32;
		cv::Scalar colour(blue, green, red); // doesn't matter in this case, but remember OpenCV uses BGR not RGB
		background = cv::Mat(1024, 1024, CV_8UC3, colour);
	}
	else
	{
		// choose one of our preloaded background images
		r = std::rand() % all_backgrounds.size();
		cv::Mat mat = all_backgrounds.at(r);

		// also apply a random rotation to the image
		r = std::rand() % 100;
		if (r <= 25)		{ background = mat.clone();										}	// no rotation -- use the original image
		else if (r <= 50)	{ cv::rotate(mat, background, cv::ROTATE_90_CLOCKWISE);			}
		else if (r <= 75)	{ cv::rotate(mat, background, cv::ROTATE_180);					}
		else				{ cv::rotate(mat, background, cv::ROTATE_90_COUNTERCLOCKWISE);	}
	}

	return background;
}


cv::Mat convert_barcode_image_to_opencv_mat(struct zint_symbol * zint)
{
	/* This next line renders the barcode as a bitmap, which populates the fields
	 * zint_symbol->bitmap, zint_symbol->bitmap_width, and zint_symbol->bitmap_height.
	 *
	 * See this blog post for details:  https://www.ccoderun.ca/programming/2019-08-16_Zint/
	 */
	ZBarcode_Buffer(zint, 0);

	cv::Mat mat(zint->bitmap_height, zint->bitmap_width, CV_8UC3);

	int bitmap_idx = 0;
	for (int y = 0; y < zint->bitmap_height; y ++)
	{
		for (int x = 0; x < zint->bitmap_width; x ++)
		{
			uint8_t * ptr = mat.ptr(y, x);
			// remember that OpenCV expects BGR and not RGB, so swap the bytes
			ptr[0] = zint->bitmap[bitmap_idx + 2];	// blue
			ptr[1] = zint->bitmap[bitmap_idx + 1];	// green
			ptr[2] = zint->bitmap[bitmap_idx + 0];	// red
			bitmap_idx += 3;
		}
	}

	return mat;
}


cv::Mat generate_random_barcode()
{
	// see this blog post for details:  https://www.ccoderun.ca/programming/2019-08-16_Zint/

	// use between 5 and 10 digits for UPC-A
	// (Zint will automatically pad with zeros as necessary)
	const size_t number_of_digits_needed = 5 + std::rand() % 6;
	std::string barcode_text;
	while (barcode_text.size() <= number_of_digits_needed)
	{
		barcode_text += std::to_string(std::rand());
	}
	// delete the extra characters to get the exact length we want
	barcode_text.erase(number_of_digits_needed);

	auto zint = ZBarcode_Create();
	zint->symbology = BARCODE_UPCA;
	ZBarcode_Encode(zint, reinterpret_cast<const uint8_t*>(barcode_text.c_str()), 0);
	cv::Mat barcode_image = convert_barcode_image_to_opencv_mat(zint);
	ZBarcode_Delete(zint);

	// randomly rotate the barcode 0, 90, 180, or 270 degrees
	cv::Mat final_result;
	const int r = std::rand() % 100;
	if (r <= 25)		{ final_result = barcode_image;												}	// no rotation
	else if (r <= 50)	{ cv::rotate(barcode_image, final_result, cv::ROTATE_90_CLOCKWISE);			}
	else if (r <= 75)	{ cv::rotate(barcode_image, final_result, cv::ROTATE_180);					}
	else				{ cv::rotate(barcode_image, final_result, cv::ROTATE_90_COUNTERCLOCKWISE);	}

	return final_result;
}


cv::Mat generate_random_background_and_barcode(const std::string & base_filename)
{
	/* This function combines the barcode image and the background image,
	 * then saves the resulting image to a .jpg file, as well as create the
	 * corresponding .txt label file needed by Darknet.
	 */

	cv::Mat image				= get_random_background_image();
	cv::Mat barcode				= generate_random_barcode();
	const int image_width		= image.cols;
	const int image_height		= image.rows;
	const int barcode_width		= barcode.cols;
	const int barcode_height	= barcode.rows;

	// find a place to insert the barcode so that it is 100% visible (not off the edge of the final image)
	int x = 0;
	int y = 0;
	while (true)
	{
		x = std::rand() % image_width;
		y = std::rand() % image_height;

		if (x + barcode_width	>= image_width	||
			y + barcode_height	>= image_height	)
		{
			// try again, need to find new X and Y coordinates
			continue;
		}

		// otherwise if we get here then we've found a good place for the barcode
		break;
	}

	/* We need to insert the barcode into the image at the (X,Y) coordinates we decided to use.  But sometimes we want
	 * to use the full barcode image (including the white background) and other times we want to use just the black
	 * lines and ignore the white background.
	 */
	cv::Mat mask;
	if (50 > std::rand() % 100)
	{
		// create the mask so we copy only the very dark bars from the barcode image, dropping out the white background
		cv::inRange(barcode, cv::Scalar(0, 0, 0), cv::Scalar(128, 128, 128), mask);
	}
	else
	{
		// create the mask so we copy the entire barcode image, including the white background
		mask = cv::Mat::ones(barcode.rows, barcode.cols, CV_8UC1);
	}

	const cv::Rect roi(cv::Point(x, y), cv::Size(barcode.cols, barcode.rows));
	barcode.copyTo(image(roi), mask);

	// save the final image to disk
	cv::imwrite(base_filename + ".jpg", image, {cv::IMWRITE_JPEG_QUALITY, 75});

	/* Create the Darknet label, which is the class index + the barcode's middle point and size.
	 * Both the middle point and the size must be normalized between 0 and 1.
	 */
	const int class_idx = 0; // in this tutorial we only have a single class ("barcode") so the class index is always zero

	const float mid_x = x + 0.5f * float(barcode_width);
	const float mid_y = y + 0.5f * float(barcode_height);

	std::ofstream ofs(base_filename + ".txt");
	ofs	<< class_idx
		<< " " << mid_x / float(image_width)
		<< " " << mid_y / float(image_height)
		<< " " << float(barcode_width)	/ float(image_width)
		<< " " << float(barcode_height)	/ float(image_height)
		<< std::endl;

	return image;
}


int main()
{
	std::srand(std::time(nullptr));

	load_all_background_images();

	const size_t number_of_files_to_create = 2000;
	std::cout << "generating " << number_of_files_to_create << " images" << std::flush;

	for (size_t idx = 0; idx < number_of_files_to_create; idx ++)
	{
		std::cout << "." << std::flush;
		generate_random_background_and_barcode("barcode_" + std::to_string(idx));
	}

	std::cout << std::endl << "Done!" << std::endl;

	return 0;
}
