#include <iostream>
#include <cstring>
#include <cassert>
#include <climits>
#include "ai_mm2.h"
#include "util.h"

using namespace std;

#define SCORE_WIN 10000
#define SCORE_LOSE (-10000)
#define SCORE_TIE 0
#define INF (1e34)

#define MAXDEPTH 3


static int scoreTable[ipow(3, RLEN)];

__attribute__((constructor))
static void initScoreTable() {
	int len = sizeof(scoreTable) / sizeof(scoreTable[0]);
	uint8_t values[RLEN];
	memset(values, 0, RLEN);

	for (int i = 0; i < len; i++) {
		int numempty = 0;
		uint8_t clr = 0;
		for (int j = 0; j < RLEN; j++) {
			if (values[j] == 0) numempty++;
			else if (clr == 0) clr = values[j];
			else if (clr != values[j]) {scoreTable[i] = 0; goto nope;}
		}

		if (numempty == RLEN) {
			scoreTable[i] = 0;
		} else {
			scoreTable[i] = (3 - 2 * clr) * (1 << (RLEN - 1 - numempty));
		}

	nope:
		for (int j = 0; j < RLEN; j++) {
			values[j]++;
			if (values[j] == 3) values[j] = 0;
			else break;
		}
	}
}


static int evaluate(Board &bd, uint8_t myclr) {
	Bounds bounds = bd.computeBounds();

	int score = 0;

#define CHECK(_ident, _xlim1, _xlim2, _ylim1, _ylim2, _dx, _dy) \
	do { \
		for (int y = (_ylim1); y <= (_ylim2); y++) { \
			for (int x = (_xlim1); x <= (_xlim2); x++) { \
				int idx = BSZ * y + x; \
				int scidx = 0; \
				for (int i = 1; i <= RLEN; i++, idx += BSZ * (_dy) + (_dx)) { \
					scidx *= 3; \
					scidx += bd[idx]; \
				} \
				score += scoreTable[scidx]; \
			} \
		} \
	} while (0)

	CHECK(1, bounds.left           , bounds.right - (RLEN-1), bounds.top, bounds.bottom           ,  1, 0);
	CHECK(2, bounds.left           , bounds.right - (RLEN-1), bounds.top, bounds.bottom - (RLEN-1),  1, 1);
	CHECK(3, bounds.left           , bounds.right           , bounds.top, bounds.bottom - (RLEN-1),  0, 1);
	CHECK(4, bounds.left + (RLEN-1), bounds.right           , bounds.top, bounds.bottom - (RLEN-1), -1, 1);

#undef CHECK

	if (myclr != 1) score = -score;
	return score;
}

static float chanceNode(
		Board &bd, uint8_t myclr, int idx, uint8_t onturn,
		int depth, float alpha, float beta);

static float playNode(
		Board &bd, uint8_t myclr, uint8_t onturn,
		int depth, float alpha, float beta) {

	if (depth == 0 || bd.bag.totalLeft() == 0) {
		return evaluate(bd, myclr);
	}

	float value;

	vector<int> edgeCells = bd.getEdgeCells();
	if (onturn == myclr) {
		value = -INF;
		for (int idx : edgeCells) {
			value = max(value, chanceNode(bd, myclr, idx, onturn, depth, alpha, beta));
			alpha = max(alpha, value);
			if (alpha >= beta) break;
		}
	} else {
		value = +INF;
		for (int idx : edgeCells) {
			value = min(value, chanceNode(bd, myclr, idx, onturn, depth, alpha, beta));
			beta = min(beta, value);
			if (alpha >= beta) break;
		}
	}

	return value;
}

static float chanceNode(
		Board &bd, uint8_t myclr, int idx, uint8_t onturn,
		int depth, float alpha, float beta) {

	const float totalLeftF = bd.bag.totalLeft();

	float score = 0;

	for (uint8_t clr = 1; clr <= NC; clr++) {
		int numLeft = bd.bag.numLeft(clr);
		if (numLeft == 0) continue;

		float probability = numLeft / totalLeftF;

		bd.bag.drawColour(clr);
		uint8_t win = bd.putCW(idx, clr);
		if (win != 0) {
			score += probability * (win == myclr ? SCORE_WIN : SCORE_LOSE);
		} else {
			score += probability * playNode(bd, myclr, NEXTTURN(onturn), depth - 1, alpha, beta);
		}

		bd.bag.replace(clr);
		bd.undo(idx);
	}

	return score;
}

int MM2::calcMove(Board &bd, uint8_t myclr) {
	assert(bd.bag.totalLeft() > 0);

	// The usage of a single score variable instead of an array works only if
	// there are only two players.
	static_assert(NC == 2, "MM only works with 2 players");

	float maxscore = INT_MIN;
	int maxat = -1;

	vector<int> edgeCells = bd.getEdgeCells();
	for (int idx : edgeCells) {
		float score = chanceNode(bd, myclr, idx, myclr, MAXDEPTH, -INF, +INF);

		if (score > maxscore) {
			maxscore = score;
			maxat = idx;
		}
	}

	return maxat;
}