/*
	ctnn, neural network harness (generate training data, self-play
	training, evaluate accuracy) using a playtak.com database

	Copyright (C) 2021, tslil clingman

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

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

#include <cnn1986.h>
#include <negamax.h>
#include <tak.h>

#define RANDPM1 ((rand()&1)?-1:+1)

enum MODE { M_RPT, M_EVL, M_GEN, M_TRN };

enum MODE mode;
FILE *training_fh = NULL;
uint64_t heights[16], samples;
float max_flats, outcome_black, loss;

// Self-play parameters
float lambda = 0.7;
int num_training_plies = 5;

const int max_depth = 6;
inline void
negamax_display_progress(const uint8_t cur_depth,
												 const uint8_t init_depth,
												 const uint32_t length) {
	(void)(cur_depth);
	(void)(init_depth);
	(void)(length);
}

static void
write_input(const int dx, const int dy, const uint8_t swap) {
	// Two numbers for flats remaining
	fprintf(training_fh,"%.8f,%.8f,",
					(float)(white_count & 127)/max_flats,
					(float)(black_count & 127)/max_flats);

	// Write the board layers
	float val;
	int col, row;
	for (uint8_t depth = 0; depth < max_depth; depth++) {
		row = (dy>0)?-1:board_size;
		for (int i = 0; i < board_size; i++) {
			row += dy;
			col = (dx>0)?-1:board_size;
			for (int j = 0; j < board_size; j++) {
				col += dx;
				const uint8_t k =
					(swap) ? THE_COORDS(row, col) : THE_COORDS(col, row);
				val = 0;
				if (COUNT_AT(k)>depth) {
					if (depth == 0) {
						// Top layer of stacks is handled differently to indicate
						// stone type
						if (STONE_AT(k) == STONE_STANDING) {
							val = (colours[k] & 1) ? +0.25 : -0.25;
						} else if (STONE_AT(k) == STONE_CAPSTONE) {
							val = (colours[k] & 1) ? +1.00 : -1.00;
						} else {
							val = (colours[k] & 1) ? +0.75 : -0.75;
						}
					} else {
						// Layers underneath
						val = (colours[k] & (1<<depth)) ? +0.75 : -0.75;
					}
				}
				fprintf(training_fh,"%.2f,", val);
			}
		}
	}
	fprintf(training_fh,"%.1f\n", outcome_black);
}

// Warning: performs _no_ checks on input whatsoever
static enum ACT_RESULT
parse_line(const char *pt, const ssize_t read) {
	ssize_t idx;
	enum ACT_RESULT r;
	int total_plies = 0;

	for (idx=0;idx<read;idx++) {
		if (pt[idx]==',') total_plies++;
	}

	const uint8_t this_game_ply = total_plies - num_training_plies;

	for(idx=0;;) {
		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': { 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;

			if (mode == M_RPT) {
				// 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;
				}
			}
		}
		switch (mode) {
			case M_RPT: break;
			case M_GEN: {
				if (ply + 3 >= total_plies)
					write_input(RANDPM1, RANDPM1, rand()&1);
				break;
			};
			case M_EVL: {
				if (ply + 3 >= total_plies) {
					samples++;
					loss += fabsf(cnn1986_evaluate_black_win() - outcome_black);
				}
				break;
			}
			case M_TRN: {
				if (ply == this_game_ply) {
					float jd[num_training_plies], grad_jd[num_training_plies][NUM_PARAMETERS];
					enum WIN_TYPE w;
					int n;

					for (n = 0; n < num_training_plies; n++) {
						jd[n] = negamax_generate();
						cnn1986_compute_gradient();
						for (int j = 0; j < NUM_PARAMETERS; j++)
							grad_jd[n][j] = cnn1986_gradient[j];

						if (do_ptn(negamax_ptn) == ACT_OK) {
							if ((w = check_win()) < 0xFF) {
								// Correct the entry
								if (w == WIN_FLAT_BLACK || w == WIN_ROAD_BLACK)
									jd[n] = +1;
								else if (w == WIN_FLAT_WHITE || w == WIN_ROAD_WHITE)
									jd[n] = -1;
								else
									jd[n] = 0;
								break;
							}
						} else {
							break;
						}
					}

					float ds[n];
					for (int k = 0; k < n; k++) ds[k] = jd[k+1] - jd[k];

				}
				break;
			}
		}

		// Parse next action
		while (idx<read && pt[idx++]!=',');
		if (idx>=read) return ACT_OK;
		next_ply();
	}
	return ACT_OK;
}

const char* license = "ctnn, neural network harness (generate training data, self-play training, evaluate accuracy) using a playtak.com database\n\
\n\
Copyright (C) 2021, tslil clingman\n\
\n\
This program comes with ABSOLUTELY NO WARRANTY; and is made available under the terms of the GNU GPL v3 license. This is free software, and you are welcome to redistribute it under certain conditions; see COPYING for details.\n";

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

	enum ACT_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,
		white_wins = 0, black_wins = 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';
	max_flats = (size == 5) ? 21.0 : 30.0;

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

	mode = M_RPT;
	if (argc > 3) {
		if (!strncmp("generate", argv[3], 8)) {
			mode = M_GEN;
			snprintf(td_fn, 64, "data/training-%d.csv", size);
			training_fh = fopen(td_fn, "w");
			if (training_fh == NULL) exit(EXIT_FAILURE);
		} else if (!strncmp("train", argv[3], 5)) {
			mode = M_TRN;
		} else if (!strncmp("evaluate", argv[3], 8)) {
			loss = 0;
			samples = 0;
			mode = M_EVL;
		}
	}

	while ((read = getline(&line, &len, playtak_fh)) != -1) {
		// Reset everything
		reset_state(size);
		// Store the outcome of this game. Black win = 0.9
		if (line[read-4] == '0') outcome_black = 0.9;
		else outcome_black = -0.9;
		// Parse the line
		r = parse_line(line,read-4);
		// Adjust counts if we're not generating training data
		if (mode == M_RPT) {
			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;
				}
				if (win == WIN_FLAT_BLACK || win == WIN_ROAD_BLACK)
					black_wins++;
				else if (win == WIN_FLAT_WHITE || win == WIN_ROAD_WHITE)
					white_wins++;
			}
		}
		games++;
	}

	fclose(playtak_fh);
	if (mode == M_GEN) fclose(training_fh);
	if (line) free(line);

	if (illegal || overflow) putchar('\n');
	printf("Read %d games\n",games);


	switch (mode) {
		case M_RPT: {
			printf("Illegals:  %d\nOverflows: %d\n\
Black wins: %.3f%%\n\
Road wins: %d\nFlat wins: %d\n\
Average turns to road win: %.3f\n\
Average turns to flat win: %.3f\n",
						 illegal, overflow,
						 (double)black_wins / (double)(black_wins+white_wins) * 100,
						 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]);
			}
			break;
		}
		case M_EVL: {
			printf("%ld samples: %.8f loss\n",
						 samples, loss / ((float)(samples)));
			break;
		}
		case M_TRN: {
			break;
		}
		case M_GEN: {
			break;
		}
	}

	exit(EXIT_SUCCESS);
}