#include <iostream>
#include <algorithm>
#include <array>
#include <climits>
#include <cassert>
#include "ai_str.h"

using namespace std;


static const int deltaxy[8][2] = {
	{1, 0}, {1, 1}, {0, 1}, {-1, 1}, {-1, 0}, {-1, -1}, {0, -1}, {1, -1}
};

static const int deltaidx[8] = {
#define ITEM(_i) (BSZ * deltaxy[_i][1] + deltaxy[_i][0])
	ITEM(0), ITEM(1), ITEM(2), ITEM(3), ITEM(4), ITEM(5), ITEM(6), ITEM(7)
#undef ITEM
};

static int findExtent(const Board &bd, int idx, int dir, uint8_t clr) {
	const int dx = deltaxy[dir][0], dy = deltaxy[dir][1], didx = deltaidx[dir];

	int x = idx % BSZ, y = idx / BSZ;
	x += dx; y += dy; idx += didx;

	for (int i = 0; i < RLEN - 1 && x > 0 && x < BSZ - 1 && y > 0 && y < BSZ - 1; i++) {
		if (bd[idx] != 0 && bd[idx] != clr) return i;
		x += dx; y += dy; idx += didx;
	}

	return RLEN - 1;
}

static int count(const Board &bd, int idx, int didx, uint8_t clr) {
	int res = 0;
	for (int i = 0; i < RLEN; i++) {
		res += bd[idx] == clr;
		idx += didx;
	}
	return res;
}

static int calcScore(const Board &bd, int idx, uint8_t myclr) {
	array<array<int, NC>, 8> exts;
	for (int dir = 0; dir < 8; dir++) {
		for (int clr = 1; clr <= NC; clr++) {
			exts[dir][clr - 1] = findExtent(bd, idx, dir, clr);
		}
	}

	array<int, NC> score;
	score.fill(0);
	for (int dir = 0; dir < 4; dir++) {
		const int dx = deltaxy[dir][0], dy = deltaxy[dir][1];
		const int didx = deltaidx[dir];
		assert(deltaxy[dir + 4][0] == -dx && deltaxy[dir + 4][1] == -dy);

		for (int clr = 1; clr <= NC; clr++) {
			int length = exts[dir][clr - 1] + exts[dir + 4][clr - 1] + 1;
			if (length < RLEN) continue;

			int start = idx + exts[dir][clr - 1] * didx;

			int maxcnt = -1;
			for (int i = 0; i < length - 4; i++) {
				int cnt = count(bd, start, -didx, clr);
				maxcnt = max(maxcnt, cnt);

				start += -didx;
			}

			// Square the count to lay emphasis on larger lines
			score[clr - 1] += maxcnt * maxcnt;
		}
	}

	int tot = 0;
	for (int i = 0; i < NC; i++) tot += score[i];

	return -tot + 3 * score[myclr - 1];
}

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

	int maxscore = INT_MIN, maxat = -1;

	vector<int> moves = bd.getEdgeCells();
	for (int idx : moves) {
		int score = calcScore(bd, idx, myclr);
		if (score > maxscore) {
			maxscore = score;
			maxat = idx;
		}
	}

	return maxat;
}