#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();
}

void
negamax_free(void) {
	zobrist_free();
}


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 WIN_TYPE w;

static float
negamax(const uint8_t cur_depth, float alpha, float beta,
				const float colour) {
	float value = -infty;

	action_list_t *list = action_list_generate();

	for (action_list_t *node=list; node!=NULL; node=node->next) {
		negamax_display_progress(cur_depth);

		action_take(node);

		// Somebody won?
		if (ply >= 2*board_size - 3 && (w = check_win()) < 0xFF) {
			if (w == WIN_ROAD_BLACK || w == WIN_FLAT_BLACK) {
				value = colour*infty;
				if (value > 0) {
					// Always take the win for ourselves
					action_undo(node);
					if (cur_depth == negamax_search_depth)
						action_to_ptn(node, negamax_ptn);
					goto prune;
				}
			} else if (w == WIN_DRAW) value = 0; // Draw is fixed at neutral
			else value = -colour*infty;
		} else if (cur_depth > 1) {
			// Recurse away from the leaves
			value = fmax(value, -negamax(cur_depth - 1, -beta, -alpha, -colour));
		} else {
			// Evaluate a leaf
			value = fmax(value, colour * cnn1986_evaluate_black_win());
		}

		action_undo(node);

		if (value > alpha) {
			alpha = value;
			if (cur_depth == negamax_search_depth)
				action_to_ptn(node, negamax_ptn);
		}
		if (alpha >= beta) break;
	}

	action_list_free(list);

 prune:
	return value;
}


/*
 * // Movement steps, orderd with the enum: UP DOWN LEFT RIGHT
 * static int8_t deltas[4];
 *
 * void negamax_init(const uint8_t new_board_size) {
 *	board_size = new_board_size;
 *	deltas[0] = +board_size;
 *	deltas[1] = -board_size;
 *	deltas[2] = -1;
 *	deltas[3] = +1;
 *	negamax_free_zobrist();
 *	negamax_init_zobrist();
 * }
 *
 * static void previous_ply(void) {
 *	if (ply>0) ply--;
 *	if (ply == 1) {
 *		current_colour = C_WHITE;
 *	} else {
 *		if (current_colour == C_BLACK) current_colour = C_WHITE;
 *		else current_colour = C_BLACK;
 *	}
 * }
 *
 * static void push_stones(const int8_t location,
 *												const uint8_t count,
 *												const uint8_t new_colours,
 *												const enum STONE_VARIANT top_stone) {
 *	colours[location]  = (colours[location] << count) | new_colours;
 *	celldat[location]  = top_stone
 *		| ((celldat[location] + ((count << NUM_SHIFT))) & NUM_MASK);
 * }
 *
 * static enum WIN_TYPE w;
 *
 * #define WIN_EVALUATE_OR_RECURSE(store,reset) {		\
 *		w = 0xFF; \
 *		if (ply >= 2*board_size - 3) w = check_win(); \
 *		if (w < 0xFF) { \
 *			/\* Somebody won, assign weights accordingly. *\/ \
 *			if (w == WIN_ROAD_BLACK || w == WIN_FLAT_BLACK) { \
 *				value = colour*infty; \
 *				/\* Always take the win *\/ \
 *				if (value > 0) { \
 *					{ reset }; \
 *					if (cur_depth == negamax_search_depth) { store }; \
 *					goto prune; \
 *				} \
 *				/\* Fix draw value to be completely neutral *\/ \
 *			} else if (w == WIN_DRAW) value = 0; \
 *			else value = -colour*infty; \
 *		} if (cur_depth == 0) { \
 *			/\* We're at the bottom, evaluate *\/ \
 *			value = fmax(value, colour * cnn1986_evaluate_black_win()); \
 *		} else { \
 *			/\* We're not at the bottom, recurse first *\/ \
 *			next_ply(); \
 *			value = fmax(value, -negamax(cur_depth - 1, -beta, -alpha, -colour)); \
 *			previous_ply(); \
 *		} \
 *		{ reset }; \
 *		/\* Update the optimal value, which alpha carries *\/ \
 *		if (value > alpha) { \
 *			alpha = value; \
 *			if (cur_depth == negamax_search_depth) { store }; \
 *			/\* Prune *\/ \
 *			if (alpha >= beta) goto prune; \
 *		} \
 *	}
 *
 * float negamax(const uint8_t cur_depth, float alpha, float beta,
 *							const float colour) {
 *
 *	/\*
 *	 * uint64_t hash = negamax_compute_zobrist();
 *	 * tt_entry_t *entry = tt_seek(hash);
 *	 * const float alpha_orig = alpha;
 *	 *
 *	 * if (entry != NULL && entry->depth >= cur_depth) {
 *	 *	if (entry->flag == TT_EXACT) {
 *	 *		return entry->value;
 *	 *	} else if (entry->flag == TT_LOWERBOUND) {
 *	 *		alpha = fmax(alpha, entry->value);
 *	 *	} else if (entry->flag == TT_UPPERBOUND) {
 *	 *		beta = fmin(beta, entry->value);
 *	 *	}
 *	 *	if (alpha >= beta) return entry->value;
 *	 * }
 *	 * enum TT_FLAG flag;
 *	 *\/
 *
 *	float value = -infty;
 *	const uint8_t black = (ply & 1),
 *		material = (black) ? black_count : white_count,
 *		flat = material & 127,
 *		cap = (ply > 2 && (material & 128)),
 *		standing = (ply > 2 && (material & 127));
 *
 *	// 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) > board_size) ? board_size : COUNT_AT(loc);
 *			// Only try moves after CPS
 *			if (count && ((colours[loc] & 1) == current_colour) && ply>2) {
 *				// There are stones, let's try moving them
 *
 *				// Pre-compute end-stops
 *				uint8_t end_stops[4][2]; // (end, not_crush)
 *				// UP DOWN LEFT RIGHT
 *				end_stops[0][0] = (board_size - row - 1 > count) ? count : board_size - row - 1;
 *				end_stops[1][0] = (row > count) ? count : row;
 *				end_stops[2][0] = (col > count) ? count : col;
 *				end_stops[3][0] = (board_size - col - 1 > count) ? count : board_size - col - 1;
 *				const uint8_t cap_top = STONE_AT(loc) == STONE_CAPSTONE;
 *				for (uint8_t d = 0; d < board_size-1; d++){
 *					end_stops[d][1] = 1;
 *					const uint8_t stop = end_stops[d][0];
 *					end_stops[d][0] = 0;
 *					for (uint8_t k = 1; k <= stop; k++) {
 *						const uint8_t stone = STONE_AT(loc+k*deltas[d]);
 *						if (stone == STONE_STANDING) {
 *							if (cap_top) {
 *								end_stops[d][1] = 0;
 *								end_stops[d][0]++;
 *							}
 *							break;
 *						} else if (stone == STONE_CAPSTONE) {
 *							break;
 *						}
 *						end_stops[d][0]++;
 *					}
 *				}
 *
 *				uint16_t colours_backup[board_size];
 *				uint8_t celldat_backup[board_size], drops[board_size];
 *				// we only ever need board_size-1 in drops actually, the
 *				// last spot is to skip a bounds check at (*)
 *
 *				//Back up the rows of the board
 *				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)];
 *				}
 *
 *				// 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 column once we start looking horizontally
 *					if (dir == M_LEFT) {
 *						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 here just try
 *					 * all the ordered partitions of num ∈ {1 … end_stop}, and
 *					 * skip the partition if it calls for multiple stones at
 *					 * the end with a crush.
 *					 *\/
 *					uint8_t gaps, t, idx, mask;
 *					for (uint8_t num = 1; num <= count; num++) {
 *						for (uint8_t steps = 1;
 *								 steps <= end_stops[dir][0] && steps <= num;
 *								 steps++) {
 *							// TODO: Generalise to board_size!
 *							gaps = 0x07 >> (board_size-steps-1);
 *							// 0b0000[0111] because 4-1=3 and 5-1=4
 *							do  {
 *								// Ensure legal move if we have to crush
 *								const uint8_t last_drop_check =
 *									(num > 1) ? (gaps & 1<<(num - 2)) : 1;
 *								if (end_stops[dir][1] || last_drop_check) {
 *									// Translate to a drop sequence
 *									drops[0] = 1; mask = 1; idx = 0;
 *									for (uint8_t d = 0; d + 1 < num; d++) {
 *										if (gaps & mask) {
 *											idx++;
 *											drops[idx] = 1; // (*) no bounds check
 *										} else {
 *											drops[idx] += 1;
 *										}
 *										mask <<= 1;
 *									}
 *									// Do it, and manually check for win if it's valid
 *									uint8_t j = num;
 *									for (uint8_t k = 0; k < steps; k++) {
 *										j -= drops[k];
 *										push_stones(loc+(k+1)*deltas[dir],
 *																drops[k],
 *																(colours[loc] >> j) & (0xFFFF >> (0x10 - drops[k])),
 *																(k == steps - 1) ? STONE_AT(loc) : STONE_FLAT);
 *									}
 *									// Then we drop them from the source
 *									colours[loc] >>= num;
 *									const uint8_t dec_count = celldat[loc] - (num << NUM_SHIFT);
 *									celldat[loc] = dec_count & NUM_MASK;
 *
 *									// First check for wins, if we're at the bottom
 *									// evaluate, otherwise recurse
 *									WIN_EVALUATE_OR_RECURSE({
 *											// If we did update the optimal value, store
 *											// this move
 *											generate_move(loc, dir, steps, drops, negamax_ptn);
 *										},{
 *											// Reset the board data after recursing or
 *											// before returning
 *											if (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];
 *												}
 *											}
 *										});
 *								}
 *								/\*
 *								 * With thanks to
 *								 * https://graphics.stanford.edu/~seander/bithacks.html#NextBitPermutation
 *								 * we have the following magic to generate the next
 *								 * permutation of steps-many set bits
 *								 *\/
 *								t = (gaps | (gaps - 1));
 *								gaps = (t + 1) | (((~t & -~t) - 1) >> (__builtin_ctz(gaps) + 1));
 *							} while (gaps && (gaps + 1 <= (1<<(num-1))));
 *						}
 *					}
 *				}
 *			} else if (material && count == 0) {
 *				// Empty square, try placements
 *
 *				if (flat) {
 *					// Generate the placement
 *					if (black) black_count--;
 *					else white_count--;
 *					colours[loc] = current_colour;
 *					celldat[loc] = NUM_INC | STONE_FLAT;
 *					WIN_EVALUATE_OR_RECURSE({
 *							// If we did update the optimal value, store
 *							generate_place(loc, STONE_FLAT, negamax_ptn);
 *						},{
 *							// Reset the state
 *							celldat[loc] = 0;
 *							if (black) black_count++;
 *							else white_count++;
 *						});
 *					// Do the same for walls, can't happen without flats
 *					if (standing) {
 *						if (black) black_count--;
 *						else white_count--;
 *						colours[loc] = current_colour;
 *						celldat[loc] = NUM_INC | STONE_STANDING;
 *						WIN_EVALUATE_OR_RECURSE({
 *								generate_place(loc, STONE_STANDING, negamax_ptn);
 *							},{
 *								celldat[loc] = 0;
 *								if (black) black_count++;
 *								else white_count++;
 *							});
 *					}
 *				}
 *
 *				// and for caps
 *				if (cap) {
 *					if (black) black_count &= 127;
 *					else white_count &= 127;
 *					colours[loc] = current_colour;
 *					celldat[loc] = NUM_INC | STONE_CAPSTONE;
 *					WIN_EVALUATE_OR_RECURSE({
 *							generate_place(loc, STONE_CAPSTONE, negamax_ptn);
 *						},{
 *							celldat[loc] = 0;
 *							if (black) black_count |= 128;
 *							else white_count |= 128;
 *						});
 *				}
 *			}
 *			negamax_display_progress(cur_depth);
 *		}
 *	}
 *
 *  prune:
 *	/\*
 *	 * flag = TT_EXACT;
 *	 * if (value <= alpha_orig) flag = TT_UPPERBOUND;
 *	 * else if (value >= beta) flag = TT_UPPERBOUND;
 *	 *
 *	 * if (entry == NULL) {
 *	 *	tt_insert(hash, flag, cur_depth, value);
 *	 * } else {
 *	 *	entry->flag = flag;
 *	 *	entry->value = value;
 *	 *	entry->depth = cur_depth;
 *	 * }
 *	 *\/
 *
 *	return value;
 * }
 */