#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <tak.h>

uint64_t heights[16];
FILE *training_fh = NULL;
int wc_col, bc_col, wc_row, bc_row, result;

static void
write_input(void) {
	// Four numbers for capstone coords (col,row)
	fprintf(training_fh,"%d,%d,%d,%d,",wc_col,wc_row,bc_col,bc_row);
	// Two numbers for flats remaining
	fprintf(training_fh,"%d,%d,",white_count & 127, black_count & 127);
	// Two layers of board_size * board_size:
	int h, t; uint16_t mask;
	for (int k = 0; k < board_size * board_size; k++) {
		// stacks encoded as balanced ternary without capstones and walls
		t = 0; mask = 1;
		h = COUNT_AT(k);
		if (STONE_AT(k)) h--;
		while (h-->0) {
			t *= 3;
			t += (colours[k] & mask) ? +1 : -1;
			mask <<= 1;
		}
		fprintf(training_fh,"%d,",t);
	}
	for (int k = 0; k < board_size * board_size; k++) {
		// walls with +- 1
		if (COUNT_AT(k) && STONE_AT(k) == STONE_STANDING)
			fprintf(training_fh,"%d,",(colours[k] & 1) ? +1 : -1);
		else fputs("0,", training_fh);
	}
}

// Warning: performs _no_ checks on input whatsoever
static enum E_RESULT
parse_line(const char *pt, const ssize_t read) {
	ssize_t idx = 0;
	enum E_RESULT r;
	for(;;) {
		if (pt[idx] == 'P') {
			// P [A-F][1-6] [CF]?,
			idx+=2;
			enum STONE_VARIANT stone;
			const uint8_t col = pt[idx]-'A', row = pt[idx+1]-'1';

			if (idx + 3 < read) {
				switch (pt[idx+3]) {
					case 'W': { stone = STONE_STANDING; break; }
					case 'C': {
						if (current_colour == C_BLACK) {
							bc_col = col + 1;
							bc_row = row + 1;
						} else {
							wc_col = col + 1;
							wc_row = row + 1;
						}
						stone = STONE_CAPSTONE;
						break;
					}
					default:  { stone = STONE_FLAT;     break; }
				}
			} else {
				stone = STONE_FLAT;
			}

			r = try_place(THE_COORDS(col,row), current_colour, stone);
			if (r != ACT_OK) return r;
		} else if (pt[idx] == 'M') {
			// M [A-F][1-6] [A-F][1-6]( [1-6])+,
			idx+=2;
			uint8_t drops[board_size];
			const uint8_t s_col =pt[idx]-'A', s_row=pt[idx+1]-'1',
				d_col=pt[idx+3]-'A', d_row=pt[idx+4]-'1';
			idx+=4;

			enum MOVE_DIRECTION dir = M_RIGHT;
			if (s_col < d_col) dir=M_RIGHT;
			else if (s_col > d_col) dir=M_LEFT;
			else if (s_row < d_row) dir=M_UP;
			else if (s_row > d_row) dir=M_DOWN;

			uint8_t steps = 0;
			do {
				idx+=2;
				drops[steps++] = pt[idx] - '0';
			} while (idx+2<read && pt[idx+1] != ',');

			r = try_move(THE_COORDS(s_col, s_row), dir, steps, drops);

			if (r != ACT_OK) return r;

			// Measure height of stacks exceeding 1
			for (int k = 0; k < board_size * board_size; k++) {
				if (COUNT_AT(k)>1) heights[COUNT_AT(k)]+=1;
			}
		}
		// Generate training data
		write_input();
		fprintf(training_fh,"%d,%d\n", result, 1-result);
		// Parse next action
		while (idx<read && pt[idx++]!=',');
		if (idx>=read) return ACT_OK;
		next_ply();
	}
	return ACT_OK;
}

int
main(int argc, char **argv) {
	(void)(argc);

	enum E_RESULT r;
	enum WIN_TYPE win;
	uint32_t games = 0, overflow=0, illegal = 0;
	uint32_t road_wins=0, flat_wins=0, road_turns=0, flat_turns=0;

	for (int k = 0; k < 16; k++) heights[k] = 0;

	size_t len = 0;
	ssize_t read = 0;
	FILE *playtak_fh = NULL;
	char *line = NULL, td_fn[65];

	const uint8_t size = argv[1][0]-'0';

	playtak_fh = fopen(argv[2], "r");
	if (playtak_fh == NULL) exit(EXIT_FAILURE);

	snprintf(td_fn, 64, "data/training-%d.csv",size);
	training_fh = fopen(td_fn, "w");
	if (training_fh == NULL) exit(EXIT_FAILURE);

	while ((read = getline(&line, &len, playtak_fh)) != -1) {
		// Reset everything
		reset_state(size);
		bc_col = 0; wc_col = 0;
		bc_row = 0; wc_row = 0;
		// Store the result of this game
		if (line[read-4] == '0') result = 1;
		else result = 0;
		// Parse the line and adjust counts
		r = parse_line(line,read-4);
		if (r == ACT_ILLEGAL) {
			illegal++;
			printf("Illegal:\n%s",line);
		} else if (r == ACT_OVERFLOW) {
			printf("Overflow:\n%s",line);
			overflow++;
		} else {
			win = check_win();
			if (win == WIN_FLAT_BLACK
					|| win == WIN_FLAT_WHITE
					|| win == WIN_DRAW) {
				flat_wins++;
				flat_turns += ply/2+1;
			} else {
				road_wins++;
				road_turns += ply/2+1;
			}
		}
		games++;
	}

	fclose(playtak_fh);
	fclose(training_fh);
	if (line) free(line);

	if (illegal || overflow) putchar('\n');
	printf("Read %d games\n",games);
	printf("Illegals:  %d\nOverflows: %d\nRoad wins: %d\nFlat wins: %d\n\
Average turns to road win: %.3f\nAverage turns to flat win: %f\n",
				 illegal,overflow, road_wins, flat_wins,
				 (double)(road_turns)/(double)(road_wins),
				 (double)(flat_turns)/(double)(flat_wins));
	for (int k = 2; k < 16; k++) {
		printf("Height %2d: %7ld\n",k,heights[k]);
	}

	exit(EXIT_SUCCESS);
}