#include "ct1975.h"

// ===================================================================
// Implementation of a small convolutional neural network
// ===================================================================

static float flattened[CONV_NUM+2];
static float dense1[DENSE1_NUM];
static float dense2[DENSE2_NUM];

#define RELU(x) ((x) = ((x)<0)?0:(x))

float
ct1975_evaluate_black_win(void) {
	/* ------------------ *
	 *  Convolution layer *
	 * ------------------ */
	// for each kernel
	for (uint8_t kern = 0; kern < KERN_NUM; kern++) {
		// the stride is 1, march across the board
		for (uint8_t bx = 0; bx < KERN_OSIZE; bx++) {
			for (uint8_t by = 0; by < KERN_OSIZE; by++) {
				flattened[kern+KERN_NUM*(bx+by*KERN_OSIZE)] =
					conv2d_biases[kern];
				// Compute the convolution for this position
				for (uint8_t ky = 0; ky < KERN_SIZE; ky++) {
					for (uint8_t kx = 0; kx < KERN_SIZE; kx++) {
						for (uint8_t c = 0; c < KERN_CHAN; c++) {
							// Where we are on the board
							const uint8_t loc = kx+by+(ky+bx)*board_size;
							// Look up what's on the board at this location, and
							// multiply it. For c=0 we have to do some extra work
							float lookup = 0;
							if (COUNT_AT(loc)>c) {
								if (c==0) {
									if (STONE_AT(loc) == STONE_STANDING) {
										lookup = (colours[loc] & 1) ? +0.25 : -0.25;
									} else if (STONE_AT(loc) == STONE_CAPSTONE) {
										lookup = (colours[loc] & 1) ? +1.00 : -1.00;
									} else {
										lookup = (colours[loc] & 1) ? +0.50 : -0.50;
									}
								} else {
									lookup = (colours[loc] & (1<<c)) ? +0.50 : -0.50;
								}
							}
							flattened[kern+KERN_NUM*(bx+by*KERN_OSIZE)]
								+= lookup*conv2d_weights[kern][ky][kx][c];
						}
					}
				}
				RELU(flattened[kern+KERN_NUM*(bx+by*KERN_OSIZE)]);
			}
		}
	}
	// Add input of flat counts
	flattened[CONV_NUM]   = (float)(white_count & 127)/21.0;
	flattened[CONV_NUM+1] = (float)(black_count & 127)/21.0;
	/* ------------------ *
	 *  First dense layer *
	 * ------------------ */
	for (uint8_t d1 = 0; d1 < DENSE1_NUM; d1++) {
		dense1[d1] = dense1_biases[d1];
		for (uint8_t fl = 0; fl < CONV_NUM+2; fl++) {
			dense1[d1] += flattened[fl]*dense1_weights[d1][fl];
		}
		RELU(dense1[d1]);
	}
	/* ------------------- *
	 *  Second dense layer *
	 * ------------------- */
	for (uint8_t d2 = 0; d2 < DENSE2_NUM; d2++) {
		dense2[d2] = dense2_biases[d2];
		for (uint8_t d1 = 0; d1 < DENSE1_NUM; d1++) {
			dense2[d2] += dense1[d1]*dense2_weights[d2][d1];
		}
		RELU(dense2[d2]);
	}
	/* ------------- *
	 *  Output layer *
	 * ------------- */
	float output = output_bias;
	for (uint8_t d2 = 0; d2 < DENSE2_NUM; d2++) {
		output += dense2[d2]*output_weights[d2];
	}

	// Truncated Pade approximant of logistic function
	output = (12.0+output+50.0*output/(output*output+10.0))/24.0;
	if (output > 1.0) return 1.0;
	else if (output < 0.0) return 0.0;

	return output;
}

// ===================================================================
// Select best move according the CNN above
// ===================================================================

char ct1975_ptn[9];
float ct1975_optimal;

#define BETTER(t,o) ((current_colour == C_BLACK && (t) > (o)) || (current_colour == C_WHITE && (t) < (o)))

void
ct1975_generate_ptn(void display_progress(void)) {
	float this = 0;
	enum E_RESULT r;
	uint16_t colours_backup[board_size];
	uint8_t celldat_backup[board_size], drops[board_size];
	const uint8_t white_count_backup = white_count,
		black_count_backup = black_count;

	// 1.0 is a `certain' black win, 0.0 is a `certain' white win.
	ct1975_optimal = (current_colour == C_BLACK) ? 0.0 : 1.0;

	// Step across the board
	for (uint8_t row = 0; row < board_size; row++) {
		for (uint8_t col = 0; col < board_size; col++) {
			// Try all valid actions for this square. Is it empty?
			const uint8_t loc = THE_COORDS(col, row);
			const uint8_t count = COUNT_AT(loc);
			// Only try moves after CPS
			if (count && ((colours[loc] & 1) == current_colour) && ply>2) {
				// There are stones, can we move them in a given direction?
				// I'm not a huge fan of looping through enums, but it's
				// better than manually unrolling this. Sufficiently smart
				// compilers?
				for (enum MOVE_DIRECTION dir = M_UP; dir <= M_RIGHT; dir++) {
					// Back-up the row/column of the board
					if (dir == M_UP || dir == M_DOWN) {
						for (uint8_t y = 0; y < board_size; y++) {
							colours_backup[y] = colours[THE_COORDS(col, y)];
							celldat_backup[y] = celldat[THE_COORDS(col, y)];
						}
					} else {
						for (uint8_t x = 0; x < board_size; x++) {
							colours_backup[x] = colours[THE_COORDS(x, row)];
							celldat_backup[x] = celldat[THE_COORDS(x, row)];
						}
					}
					// We don't do anything terribly efficient or smart here,
					// just try everything...

					// For every number of steps
					for (uint8_t steps = 1; steps < board_size && steps <= count; steps++) {
						uint8_t idx, carry;
						for (idx = 0; idx < steps; idx++) drops[idx]=0;
						idx = 0;
						while (idx < steps) {
							// Increment the drop sequence
							carry = 0;
							drops[idx]++;
							do {
								if (carry) { drops[++idx]++; carry = 0;}
								if (drops[idx] > count || drops[idx] > board_size) {
									drops[idx] = 1; carry = 1;
								}
							} while (carry && idx < steps);
							// If carry is still set here we're done
							if (carry == 0) {
								// Try it, and note that try_move will never return
								// GAME_END. It does not check for winners. We don't
								// presently do that either, trust in the magic
								// numbers :)
								r = try_move(loc, dir, steps, drops);
								if (r == ACT_OK) {
									this = ct1975_evaluate_black_win();
									if (BETTER(this, ct1975_optimal)) {
										// Update the chosen action
										ct1975_optimal = this;
										generate_move(board_size, loc, dir, steps, drops, ct1975_ptn);
									}
									// Reset the board data
									if (dir == M_UP || dir == M_DOWN) {
										for (uint8_t y = 0; y < board_size; y++) {
											colours[THE_COORDS(col, y)] = colours_backup[y];
											celldat[THE_COORDS(col, y)] = celldat_backup[y];
										}
									} else {
										for (uint8_t x = 0; x < board_size; x++) {
											colours[THE_COORDS(x, row)] = colours_backup[x];
											celldat[THE_COORDS(x, row)] = celldat_backup[x];
										}
									}
								}
							}
						}
					}
				}
			} else if (count == 0) {
				// Empty square, try the three placements. Again, looping
				// through enums, sigh.
				for (enum STONE_VARIANT stone = STONE_FLAT;
						 stone <= STONE_CAPSTONE; stone++) {
					// try_place will never check for winning, and we don't do
					// that either here
					r = try_place(loc, current_colour, stone);
					// Legal placement, evaluate it
					if (r == ACT_OK) {
						this = ct1975_evaluate_black_win();
						if (BETTER(this, ct1975_optimal)) {
							// Update the chosen action
							ct1975_optimal = this;
							generate_place(board_size, loc, stone, ct1975_ptn);
						}
						// Reset the state
						celldat[loc] = 0;
						white_count = white_count_backup;
						black_count = black_count_backup;
					}
				}
			}
			display_progress();
		}
	}
}