#include "gen.h"
#include <cassert>

/**
 * Author: Niyaz Nigmatullin
 */

struct cell {
	int row;
	int col;

	cell() : row(), col() {}
	cell(int row, int col) : row(row), col(col) {}
};

vector<int> gen_path_straight(int size) {
	int r = 0;
	int c = 0;
	vector<int> res;
	while (r + 1 < size || c + 1 < size) {
		int dir;
		if (c + 1 >= size) {
			dir = 3;
		} else if (r + 1 >= size) {
			dir = 1;
		} else if (rnd.next(2) == 0) {
			dir = 3;
		} else {
			dir = 1;
		}
		r += DR[dir];
		c += DC[dir];
		res.push_back(dir);
	}
	return res;
}

vector<int> gen_random_cube() {
	vector<int> cube(6);
	for (int i = 0; i < 6; i++) {
		cube[i] = rnd.next(MAX_COLORS);
	}
	return cube;
}

vector<vector<int>> gen_straight_field(int size) {
	vector<vector<int>> res(size, vector<int>(size, -1));
	vector<int> path = gen_path_straight(size);
	vector<int> cube = gen_random_cube();
	{
		int r = 0;
		int c = 0;
		for (int dir: path) {
			res[r][c] = cube[0];
			cube = rotate(std::move(cube), dir);
			r += DR[dir];
			c += DC[dir];
		}
		res[r][c] = cube[0];
	}
	for (int r = 0; r < size; r++) {
		for (int c = 0; c < size; c++) {
			if (res[r][c] == -1) {
				if (rnd.next(2) == 0) {
					res[r][c] = cube[rnd.next(6)];
				} else {
					res[r][c] = rnd.next(MAX_COLORS);
				}
			}
		}
	}
	return res;
}

vector<int> gen_color_set(int all_count, int set_count) {
	vector<int> used(set_count, 0);
	int left_to_use = set_count;
	vector<int> res(all_count);
	for (int i = 0; i < all_count; i++) {
		bool need_to_choose_unused = left_to_use == all_count - i;
		int color = -1;
		if (need_to_choose_unused) {
			int which = rnd.next(left_to_use);
			for (int j = 0; j < set_count; j++) {
				if (!used[j]) {
					if (which == 0) {
						color = j;
						break;
					}
					--which;
				}
			}
			assert(color >= 0);
		} else {
			color = rnd.next(set_count);
		}
		if (!used[color]) {
			left_to_use--;
			used[color] = 1;
		}
		res[i] = color;
	}
	return res;
}

vector<vector<int>> gen_cube_field(int size, int path_colors) {
	vector<vector<int>> field(size, vector<int>(size));
	assert(1 <= path_colors && path_colors <= 4);
	vector<int> path = gen_color_set(4, path_colors);
	for (int i = 0; i < size; i += 2) {
		for (int j = 0; j < size; j++) {
			bool alter = (i & 2) != 0;
			int in_path = alter ? ((2 - j) & 3) : (j & 3);
			field[i][j] = path[in_path];
		}
	}
	for (int i = 1; i < size; i += 4) {
		for (int j = 0; j < size; j++) {
			field[i][j] = j + size;
		}
	}
	int block = rnd.next((size - 3) / 4);
	for (int i = 3; i < size; i += 4) {
		for (int j = 0; j < size; j++) {
			if (i == block * 4 + 3) {
				field[i][j] = 4 * size;
			} else {
				field[i][j] = j + 2 * size;
			}
		}
		if (opt<int>("no_block", 0) && i == block * 4 + 3) {
			int j = rnd.next(size);
			field[i][j] = j + 2 * size;
		}
	}
	return field;
}

vector<string> const TYPES = {"random", "straight", "cube"};

int main(int argc, char **argv) {
	registerGen(argc, argv, 1);
	int total_cells = 0;
	vector<test> tests;
	int n = opt<int>("n");
	int t = opt<int>("t", MAX_TOTAL_CELLS);
	int type = (int) (std::find(TYPES.begin(), TYPES.end(), opt<string>("type", "straight")) - TYPES.begin());
	assert(type < (int) TYPES.size());
	while ((int) tests.size() < t && total_cells + n * n <= MAX_TOTAL_CELLS) {
		int cur_type = type == 0 ? rnd.next((int) TYPES.size() - 1) + 1 : type;
		if (cur_type == 1) {
			tests.push_back(test { n, gen_straight_field(n) });
		} else if (cur_type == 2) {
			int path_colors = has_opt("path_colors") ? opt<int>("path_colors") : rnd.next(1, 4);
			tests.push_back(test { n, gen_cube_field(n, path_colors) });
		} else {
			assert(false);
		}
		total_cells += n * n;
	}
	printTests(tests);
}
