#include "negamax.h"

// ===================================================================
// Globals
// ===================================================================

const float infty = 3.0;
char negamax_ptn[9];
uint8_t negamax_search_depth = 3;

static uint64_t *zobrist[15];

// ===================================================================
// Helpers
// ===================================================================

static void
zobrist_free(void);

static int
zobrist_init(void);

static uint64_t
zobrist_compute(void);

static float
negamax(const uint8_t cur_depth, float alpha, float beta,
				const float colour);

// ===================================================================
// Zobrist hashing
// ===================================================================

static uint64_t
zobrist_compute(void) {
	uint64_t hash = 0;
	for (uint8_t l=0; l<board_size*board_size; l++) {
		colour_stack_t c = colours[l];
		const uint8_t count = COUNT_AT(l);
		enum STONE_VARIANT s = STONE_AT(l);
		for (uint8_t h=0; h<15; h++) {
			if (h<count) {
				hash ^= zobrist[h][l*(2*3+1)+(c&1)*3+s];
				c >>= 1;
			}
		}
	}
	return hash;
}

static int
zobrist_init(void) {
	for (int k=0; k<15; k++) {
		if (zobrist[k] != NULL) return EXIT_FAILURE;
	}

	for (int j=0; j<15; j++) {
		zobrist[j] = malloc(sizeof(uint64_t)*board_size*board_size*(2*3+1));
		for (int k=0; k<board_size*board_size*(2*3+1); k++) {
			XORSHIFT64;
			zobrist[j][k] = RANDOM64;
		}
	}

	return EXIT_SUCCESS;
}

static void
zobrist_free(void) {
	for (int k=0; k<15; k++) {
		if (zobrist[k] != NULL) {
			free(zobrist[k]);
			zobrist[k] = NULL;
		}
	}
}

// ===================================================================
// α-β negamax using the cnn1986 evaluation function and transposition
// tables using Zobrist hasing and a treap
// ===================================================================

void
negamax_init(const uint8_t new_board_size) {
	board_size = new_board_size;
	action_list_init();
	zobrist_init();
	tt_init();
}

void
negamax_free(void) {
	zobrist_free();
}

static enum WIN_TYPE w;

static inline float
negamax_evaluate_terminal(const float colour) {
	if (ply >= 2*board_size - 2 && (w = check_win()) < 0xFF) {
		// Check win if far enough into the game
		if (w == WIN_ROAD_BLACK || w == WIN_FLAT_BLACK) return colour*infty;
		else if (w == WIN_DRAW) return 0; // Draw is fixed at neutral
		else return -colour*infty;
	} else {
		return colour * cnn1986_evaluate_black_win();
	}
}

float
negamax_generate(void) {
	// We need to start with something outside of [-∞,∞] because those
	// values are wins
	const float safe_infty = infty + 1;

	tt_init();
	float result = negamax(negamax_search_depth,
												 -safe_infty, safe_infty,
												 (ply & 1) ? +1.0 : -1.0);
	tt_free();

	return result;
}

static enum TT_FLAG flag;

static float
negamax(const uint8_t cur_depth, float alpha, float beta,
				const float colour) {

	uint64_t hash = zobrist_compute();
	tt_entry_t *entry = tt_seek(hash);

	// CAUTION: >= breaks search stability
	if (entry != NULL && entry->depth == cur_depth) {
		if (entry->flag == TT_EXACT) {
			return entry->value;
		} else if (entry->flag == TT_LOWERBOUND && entry->value > alpha) {
			alpha = entry->value;
		} else if (entry->flag == TT_UPPERBOUND && entry->value < beta) {
			beta = entry->value;
		}
		if (alpha >= beta) return entry->value;
	}

	action_list_t *list;
	if ((list = action_list_generate()) == NULL)
		return alpha; // should never happen!

	if (entry != NULL) {
		action_move_to_front(entry->action, list);
	}

	action_t best_action;
	float value = -infty;

	for (action_node_t *node=list->head; node!=NULL; node=node->next) {

		action_take(node->action);
		// Compute the value of the node
		float node_value;
		if (ply >= 2*board_size - 2 && (w = check_win()) < 0xFF) {
			node_value = -colour*infty;
			// Check win if far enough into the game
			if (w == WIN_ROAD_BLACK || w == WIN_FLAT_BLACK) node_value = colour*infty;
			else if (w == WIN_DRAW) node_value = 0; // Draw is neutral
		} else if (cur_depth > 1) {
			// If nobody won, or too early and not leaf, recurse
			node_value = -negamax(cur_depth - 1, -beta, -alpha, -colour);
		} else {
			node_value = colour * cnn1986_evaluate_black_win();
		}
		action_undo(node->action);

		negamax_display_progress(cur_depth, list->length);

		if (node_value > value) {
			value = node_value;
			best_action = node->action;
			if (cur_depth == negamax_search_depth)
				action_to_ptn(node->action, negamax_ptn);
		}

		alpha = fmax(value, alpha);
		if (alpha >= beta) break;
	}

	action_list_free(list);

	flag = TT_EXACT;
	if (value >= beta) flag = TT_LOWERBOUND;
	else if (value <= alpha) flag = TT_UPPERBOUND;

	if (entry == NULL) {
		tt_insert(hash, flag, cur_depth, value, best_action);
	} else {
		entry->flag = flag;
		entry->value = value;
		entry->depth = cur_depth;
		entry->action = best_action;
	}

	return value;
}