#include <cstdlib>
#include <cassert>
#include "board.h"


static inline int pawnShiftDeltaY(int clr) {
	return clr == WHITE ? -1 : 1;
}

static inline int homerowY(int clr) {
	return clr == WHITE ? 7 : 0;
}


Board::Board() {
	bd[0] = -ROOK; bd[1] = -KNIGHT; bd[2] = -BISHOP; bd[3] = -QUEEN;
	bd[4] = -KING; bd[5] = -BISHOP; bd[6] = -KNIGHT; bd[7] = -ROOK;
	for (int i = 8; i < 16; i++) bd[i] = -PAWN;

	memset(bd + 16, 0, 32 * sizeof(int));

	for (int i = 48; i < 56; i++) bd[i] = PAWN;
	bd[56] = ROOK; bd[57] = KNIGHT; bd[58] = BISHOP; bd[59] = QUEEN;
	bd[60] = KING; bd[61] = BISHOP; bd[62] = KNIGHT; bd[63] = ROOK;

	pieceScore[0] = pieceScore[1] = 0;
	lastTarget = -1;
	onTurn = WHITE;
	kingMoved[0] = kingMoved[1] = false;
	rookLmoved[0] = rookLmoved[1] = false;
	rookRmoved[0] = rookRmoved[1] = false;
}

int Board::locate(int stone) const {
	for (int i = 0; i < 64; i++) {
		if (at(i) == stone) return i;
	}
	return -1;
}

#define SHIFTER(from_, to_) \
	[f_ = (from_), t_ = (to_)](Board &B) { \
		B.pieceScore[colour(B.at(f_))] += worth(B.at(t_)); \
		B.at(t_) = B.at(f_); B.at(f_) = EMPTY; B.lastTarget = t_; \
		B.onTurn = !B.onTurn; \
	}

// ~390k runs per second
vector<Board> Board::subsequents() const {
	vector<Board> subs;

	for (int i = 0; i < 64; i++) {
		int stone = at(i), clr = colour(stone);
		if (clr != onTurn) continue;
		int x = i % 8, y = i / 8;

		switch (abs(stone)) {
			case EMPTY:
				break;

			case PAWN: {
				int delta = pawnShiftDeltaY(clr);
				int homerow = homerowY(clr) + delta;
				int lastRow = homerowY(!clr);
				if (y == homerow && at(x, y + delta) == EMPTY && at(x, y + 2 * delta) == EMPTY) {
					subs.push_back(newWith(SHIFTER(i, IDX(x, y + 2 * delta))));
				}
				if (y != lastRow && at(x, y + delta) == EMPTY) {
					subs.push_back(newWith(SHIFTER(i, IDX(x, y + delta))));
				}
				if (y != lastRow && x != 0 && at(x - 1, y + delta) != EMPTY &&
						colour(at(x - 1, y + delta)) != clr) {
					subs.push_back(newWith(SHIFTER(i, IDX(x - 1, y + delta))));
				}
				if (y != lastRow && x != 7 && at(x + 1, y + delta) != EMPTY &&
						colour(at(x + 1, y + delta)) != clr) {
					subs.push_back(newWith(SHIFTER(i, IDX(x + 1, y + delta))));
				}
				if (x != 0 && y == homerow + 3 * delta &&
						at(x - 1, y) == -stone && lastTarget == IDX(x - 1, y)) {
					subs.push_back(newWith([x, y, delta, clr](Board &B) {
						B.pieceScore[clr] += worth(PAWN);
						B.at(x - 1, y + delta) = B.at(x, y);
						B.at(x, y) = EMPTY;
						B.at(x - 1, y) = EMPTY;
						B.lastTarget = IDX(x - 1, y + delta);
						B.onTurn = !B.onTurn;
					}));
				}
				if (x != 7 && y == homerow + 3 * delta &&
						at(x + 1, y) == -stone && lastTarget == IDX(x + 1, y)) {
					subs.push_back(newWith([x, y, delta, clr](Board &B) {
						B.pieceScore[clr] += worth(PAWN);
						B.at(x + 1, y + delta) = B.at(x, y);
						B.at(x, y) = EMPTY;
						B.at(x + 1, y) = EMPTY;
						B.lastTarget = IDX(x + 1, y + delta);
						B.onTurn = !B.onTurn;
					}));
				}
				break;
			}

			case ROOK: {
				int homerow = homerowY(clr);
				int which = y == homerow ? (x == 0 ? -1 : x == 7) : 0;
				for (int x2 = x - 1; x2 >= 0; x2--) {
					if (at(x2, y) == EMPTY || colour(at(x2, y)) != clr) {
						subs.push_back(newWith([i, x, y, x2, which, clr](Board &B) {
							SHIFTER(i, IDX(x2, y))(B);
							if (which == -1) B.rookLmoved[clr] = true;
							else if (which == 1) B.rookRmoved[clr] = true;
						}));
					}
					if (at(x2, y) != EMPTY) break;
				}
				for (int x2 = x + 1; x2 < 8; x2++) {
					if (at(x2, y) == EMPTY || colour(at(x2, y)) != clr) {
						subs.push_back(newWith([i, x, y, x2, which, clr](Board &B) {
							SHIFTER(i, IDX(x2, y))(B);
							if (which == -1) B.rookLmoved[clr] = true;
							else if (which == 1) B.rookRmoved[clr] = true;
						}));
					}
					if (at(x2, y) != EMPTY) break;
				}
				for (int y2 = y - 1; y2 >= 0; y2--) {
					if (at(x, y2) == EMPTY || colour(at(x, y2)) != clr) {
						subs.push_back(newWith([i, x, y, y2, which, clr](Board &B) {
							SHIFTER(i, IDX(x, y2))(B);
							if (which == -1) B.rookLmoved[clr] = true;
							else if (which == 1) B.rookRmoved[clr] = true;
						}));
					}
					if (at(x, y2) != EMPTY) break;
				}
				for (int y2 = y + 1; y2 < 8; y2++) {
					if (at(x, y2) == EMPTY || colour(at(x, y2)) != clr) {
						subs.push_back(newWith([i, x, y, y2, which, clr](Board &B) {
							SHIFTER(i, IDX(x, y2))(B);
							if (which == -1) B.rookLmoved[clr] = true;
							else if (which == 1) B.rookRmoved[clr] = true;
						}));
					}
					if (at(x, y2) != EMPTY) break;
				}
				break;
			}

			case KNIGHT:
				for (int dy = -1; dy <= 1; dy += 2)
				for (int dx = -1; dx <= 1; dx += 2)
				for (int onetwo = 1; onetwo <= 2; onetwo++) {
					int x2 = x + onetwo * dx, y2 = y + (3 - onetwo) * dy;
					if (x2 >= 0 && x2 < 8 && y2 >= 0 && y2 < 8 &&
							(at(x2, y2) == EMPTY || colour(at(x2, y2)) != clr)) {
						subs.push_back(newWith(SHIFTER(i, IDX(x2, y2))));
					}
				}
				break;

			case BISHOP:
				if (i % 8 != 0) {
					for (int i2 = i - 9; i2 >= 0; i2 -= 9) {
						if (at(i2) == EMPTY || colour(at(i2)) != clr) {
							subs.push_back(newWith(SHIFTER(i, i2)));
						}
						if (at(i2) != EMPTY) break;
						if (i2 % 8 == 0) break;
					}
				}
				if (i % 8 != 7) {
					for (int i2 = i - 7; i2 >= 0; i2 -= 7) {
						if (at(i2) == EMPTY || colour(at(i2)) != clr) {
							subs.push_back(newWith(SHIFTER(i, i2)));
						}
						if (at(i2) != EMPTY) break;
						if (i2 % 8 == 7) break;
					}
				}
				if (i % 8 != 0) {
					for (int i2 = i + 7; i2 < 64; i2 += 7) {
						if (at(i2) == EMPTY || colour(at(i2)) != clr) {
							subs.push_back(newWith(SHIFTER(i, i2)));
						}
						if (at(i2) != EMPTY) break;
						if (i2 % 8 == 0) break;
					}
				}
				if (i % 8 != 7) {
					for (int i2 = i + 9; i2 < 64; i2 += 9) {
						if (at(i2) == EMPTY || colour(at(i2)) != clr) {
							subs.push_back(newWith(SHIFTER(i, i2)));
						}
						if (at(i2) != EMPTY) break;
						if (i2 % 8 == 7) break;
					}
				}
				break;

			case KING: {
				for (int dy = -1; dy <= 1; dy++)
				for (int dx = -1; dx <= 1; dx++) {
					if (dx == 0 && dy == 0) continue;
					int x2 = x + dx, y2 = y + dy;
					if (x2 < 0 || x2 >= 8 || y2 < 0 || y2 >= 8) continue;
					if (at(x2, y2) == EMPTY || colour(at(x2, y2)) != clr) {
						subs.push_back(newWith([i, x2, y2, clr](Board &B) {
							SHIFTER(i, IDX(x2, y2))(B);
							B.kingMoved[clr] = true;
						}));
					}
				}
				int homerow = homerowY(clr);
				if (!kingMoved[clr]) {
					if (!rookLmoved[clr] &&
							at(1, homerow) == EMPTY && at(2, homerow) == EMPTY) {
						subs.push_back(newWith([i, homerow, clr](Board &B) {
							B.at(i - 1) = B.at(i - 3);
							B.at(i - 3) = EMPTY;
							B.at(i - 2) = B.at(i);
							B.at(i - i) = EMPTY;
							B.kingMoved[clr] = true;
							B.rookLmoved[clr] = true;
							B.onTurn = !B.onTurn;
						}));
					}
					if (!rookRmoved[clr] &&
							at(4, homerow) == EMPTY && at(5, homerow) == EMPTY && at(6, homerow) == EMPTY) {
						subs.push_back(newWith([i, homerow, clr](Board &B) {
							B.at(i + 2) = B.at(i + 4);
							B.at(i + 4) = EMPTY;
							B.at(i + 3) = B.at(i);
							B.at(i) = EMPTY;
							B.kingMoved[clr] = true;
							B.rookRmoved[clr] = true;
							B.onTurn = !B.onTurn;
						}));
					}
				}
				break;
			}

			case QUEEN:
				for (int dy = -1; dy <= 1; dy++)
				for (int dx = -1; dx <= 1; dx++) {
					if (dx == 0 && dy == 0) continue;
					for (int k = 1; k < 8; k++) {
						int x2 = x + k * dx, y2 = y + k * dy;
						if (x2 < 0 || x2 >= 8 || y2 < 0 || y2 >= 8) break;
						if (at(x2, y2) == EMPTY || colour(at(x2, y2)) != clr) {
							subs.push_back(newWith(SHIFTER(i, IDX(x2, y2))));
						}
						if (at(x2, y2) != EMPTY) break;
					}
				}
				break;
		}
	}

	return subs;
}

void Board::apply(const Move &mv) {
	assert(mv.castle == 0);
	onTurn = !onTurn;
	int dx = mv.to % 8 - mv.from % 8;
	if (abs(dx) == 1 && lastTarget == mv.from + dx) {
		if (at(mv.from) == PAWN && at(mv.from + dx) == -PAWN &&
				mv.from / 8 == homerowY(BLACK) + 3 * pawnShiftDeltaY(BLACK) &&
				mv.to / 8 == homerowY(BLACK) + 2 * pawnShiftDeltaY(BLACK)) {
			at(mv.from + dx) = EMPTY;
			at(mv.to) = PAWN;
			at(mv.from) = EMPTY;
			lastTarget = mv.to;
			pieceScore[WHITE] += worth(PAWN);
			return;
		} else if (at(mv.from) == -PAWN && at(mv.from + dx) == PAWN &&
				mv.from / 8 == homerowY(WHITE) + 3 * pawnShiftDeltaY(WHITE) &&
				mv.to / 8 == homerowY(WHITE) + 2 * pawnShiftDeltaY(WHITE)) {
			at(mv.from + dx) = EMPTY;
			at(mv.to) = -PAWN;
			at(mv.from) = EMPTY;
			lastTarget = mv.to;
			pieceScore[BLACK] += worth(PAWN);
			return;
		}
	}
	int clr = colour(at(mv.from));
	int homerow = homerowY(clr);
	int x = mv.from % 8, y = mv.from / 8;
	if (y == homerow && abs(at(mv.from)) == ROOK) {
		if (x == 0) rookLmoved[clr] = true;
		else if (x == 7) rookRmoved[clr] = true;
	} else if (abs(at(mv.from)) == KING) {
		kingMoved[clr] = true;
	}
	pieceScore[clr] += worth(at(mv.to));
	at(mv.to) = at(mv.from);
	at(mv.from) = EMPTY;
	lastTarget = mv.to;
}

bool Board::isValid(const Move &mv) const {
	assert(mv.castle == 0);
	if (colour(at(mv.from)) != onTurn) return false;
	Board B(*this);
	B.apply(mv);
	bool found = false;
	for (const Board &B2 : subsequents()) {
		if (B == B2) {
			found = true;
			break;
		}
	}
	if (!found) return false;
	switch (B.checkCheck(onTurn)) {
		case CheckResult::check:
		case CheckResult::checkmate:
			return false;
		default:
			return true;
	}
}

CheckResult Board::checkCheck(int player) const {
	// First check 'check' (p), then check if we have 'check' after every move (q).
	// p | q | Result
	// 0 | 0 | safe
	// 0 | 1 | stalemate
	// 1 | 0 | check
	// 1 | 1 | checkmate


	// Check 'check' by switching onTurn and seeing if the king can be taken
	bool haveP = false;
	Board B1(*this);
	B1.onTurn = !player;
	for (const Board &B2 : B1.subsequents()) {
		if (B2.pieceScore[B1.onTurn] >= worth(KING)) {
			haveP = true;
			break;
		}
	}

	// Then check if we have 'check' after every move
	bool haveQ = true;
	B1.onTurn = player;
	for (const Board &B2 : subsequents()) {
		bool checkNow = false;
		for (const Board &B3 : B2.subsequents()) {
			if (B3.pieceScore[B2.onTurn] >= worth(KING)) {
				checkNow = true;
				break;
			}
		}
		if (!checkNow) {
			haveQ = false;
			break;
		}
	}

	if (haveP) {
		if (haveQ) return CheckResult::checkmate;
		else return CheckResult::check;
	} else {
		if (haveQ) return CheckResult::stalemate;
		else return CheckResult::safe;
	}
}

static const char* playerName(int player) {
	return &"WHITE\0BLACK"[player * 6];
}

ostream& operator<<(ostream &os, const Board &B) {
	for (int y = 0; y < 8; y++) {
		os << "\x1B[34m" << "87654321"[y] << "\x1B[0m ";
		for (int x = 0; x < 8; x++) {
			if (x != 0) os << ' ';
			if (B.at(x, y) >= 0) {
				os << ".PRNBKQ"[B.at(x, y)];
			} else {
				os << ".prnbkq"[-B.at(x, y)];
			}
		}
		if (y == 0) {
			os << "  ps[W] = " << B.pieceScore[WHITE];
			os << "  onTurn = " << playerName(B.onTurn);
		} else if (y == 1) {
			os << "  ps[B] = " << B.pieceScore[BLACK];
		} else if (y == 2) {
			switch (B.checkCheck(B.onTurn)) {
				case CheckResult::safe: break;
				case CheckResult::check: os << "  " << playerName(B.onTurn) << " IN CHECK"; break;
				case CheckResult::checkmate: os << "  " << playerName(B.onTurn) << " CHECKMATED."; break;
				case CheckResult::stalemate: os << "  " << playerName(B.onTurn) << " STALEMATED."; break;
			}
		}
		os << '\n';
	}
	os << "  \x1B[34ma b c d e f g h\x1B[0m" << '\n';
	return os;
}

bool operator==(const Board &B1, const Board &B2) {
	bool c1 = memcmp(B1.bd, B2.bd, 64 * sizeof B1.bd[0]) == 0;
	bool c2 = B1.lastTarget == B2.lastTarget;
	bool c3 = B1.onTurn == B2.onTurn;
	bool c4 = memcmp(B1.kingMoved, B2.kingMoved, 2 * sizeof(bool)) == 0;
	bool c5 = memcmp(B1.rookLmoved, B2.rookLmoved, 2 * sizeof(bool)) == 0;
	bool c6 = memcmp(B1.rookRmoved, B2.rookRmoved, 2 * sizeof(bool)) == 0;
	return c1 && c2 && c3 && c4 && c5 && c6;
}