/*
	This file is part of ct.

	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 ct. If not, see <https://www.gnu.org/licenses/>.
*/

#include "cnn1986.h"
#include "weights.h"

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

static float cur_board[5*5][KERN_CHAN];
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 cnn1986_evaluate_black_win(void) {
	/* --------------- *
	 *  Populate input *
	 * --------------- */
	for (unsigned int y = 0; y < board_size; y++) {
		for (unsigned int x = 0; x < board_size; x++) {
			const unsigned int loc = x+y*board_size;
			const unsigned int count = COUNT_AT(loc);
			colour_stack_t colour = colours[loc];

			if (count > 0) {
				float lookup = 0;
				if (STONE_AT(loc) == STONE_STANDING) {
					lookup = (colour & 1) ? +0.25 : -0.25;
				} else if (STONE_AT(loc) == STONE_CAPSTONE) {
					lookup = (colour & 1) ? +1.00 : -1.00;
				} else {
					lookup = (colour & 1) ? +0.50 : -0.50;
				}
				cur_board[loc][0] = lookup;

				colour>>=1;
				for (unsigned int c = 1; c < count && c < KERN_CHAN; c++, colour>>=1) {
					cur_board[loc][c] = (colour & 1) ? +0.50 : -0.50;
				}

				for (unsigned int c = count; c < KERN_CHAN; c++) {
					cur_board[loc][c] = 0;
				}
			} else {
				for (unsigned int c = 0; c < KERN_CHAN; c++) {
					cur_board[loc][c] = 0;
				}
			}
		}
	}
	/* ------------------ *
	 *  Convolution layer *
	 * ------------------ */
	// for each kernel
	for (unsigned int kern = 0; kern < KERN_NUM; kern++) {
		// the stride is 1, march across the board
		for (unsigned int bx = 0; bx < KERN_OSIZE; bx++) {
			for (unsigned int by = 0; by < KERN_OSIZE; by++) {
				flattened[kern+KERN_NUM*(bx+by*KERN_OSIZE)] =
					conv2d_biases[kern];
				// Compute the convolution for this position
				for (unsigned int ky = 0; ky < KERN_SIZE; ky++) {
					for (unsigned int kx = 0; kx < KERN_SIZE; kx++) {
						for (unsigned int c = 0; c < KERN_CHAN; c++) {
							// Where we are on the board
							const unsigned int loc = kx+bx+(ky+by)*board_size;
							flattened[kern+KERN_NUM*(bx+by*KERN_OSIZE)]
								+= cur_board[loc][c]*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 (unsigned int d1 = 0; d1 < DENSE1_NUM; d1++) {
		dense1[d1] = dense1_biases[d1];
		for (unsigned int fl = 0; fl < CONV_NUM+2; fl++) {
			dense1[d1] += flattened[fl]*dense1_weights[d1][fl];
		}
		RELU(dense1[d1]);
	}
	/* ------------------- *
	 *  Second dense layer *
	 * ------------------- */
	for (unsigned int d2 = 0; d2 < DENSE2_NUM; d2++) {
		dense2[d2] = dense2_biases[d2];
		for (unsigned int d1 = 0; d1 < DENSE1_NUM; d1++) {
			dense2[d2] += dense1[d1]*dense2_weights[d2][d1];
		}
		RELU(dense2[d2]);
	}
	/* ------------- *
	 *  Output layer *
	 * ------------- */
	float output = output_bias;
	for (unsigned int d2 = 0; d2 < DENSE2_NUM; d2++) {
		output += dense2[d2]*output_weights[d2];
	}
	// 2*(clamped Pade approximant of logistic function) - 1
	output = (12.0+output+50.0*output/(output*output+10.0))/12.0 - 1.0;
	if (output > 1.0) {
		return 1.0;
	}
	else if (output < -1.0) {
		return -1.0;
	}
	return output;
}