from collections import Counter

import numpy as np
from draughts import RussianBoard
from draughts.models import Color
from loguru import logger

logger.remove()
NAMES = RussianBoard.SQUARE_NAMES


def encode(board, perspective):
    masks = (board.white_men, board.white_kings, board.black_men, board.black_kings)
    if perspective == Color.BLACK:
        masks = masks[2:] + masks[:2]
    x = np.zeros(128, np.float32)
    for p, mask in enumerate(masks):
        while mask:
            bit = mask & -mask
            s = bit.bit_length() - 1
            x[p * 32 + (s if perspective == Color.WHITE else 31 - s)] = 1
            mask ^= bit
    return x


def key(board):
    return (board.white_men, board.white_kings, board.black_men, board.black_kings, board.turn.value)


def algebraic_fen(board):
    sides = []
    for label, man, king in (("W", board.white_men, board.white_kings), ("B", board.black_men, board.black_kings)):
        fields = [("K" if king & (1 << s) else "") + NAMES[s] for s in range(32) if (man | king) & (1 << s)]
        sides.append(label + ",".join(fields))
    return ("W" if board.turn == Color.WHITE else "B") + ":" + ":".join(sides)


def move_text(move):
    return (":" if move.captured_list else "-").join(NAMES[s] for s in move.square_list)


class Adjudicator:
    def __init__(self, board, max_plies=300):
        self.max_plies = max_plies
        self.plies = 0
        self.seen = Counter({key(board): 1})
        self.quiet_kings = 0
        self.material_stable = 0
        self.regulation_key = None
        self.regulation_moves = [0, 0]
        self._regulation(board)

    def _regulation(self, b):
        counts = (b.white_men.bit_count(), b.white_kings.bit_count(), b.black_men.bit_count(), b.black_kings.bit_count())
        wm, wk, bm, bk = counts
        lone_side = 0 if wm == 0 and wk == 1 else 1 if bm == 0 and bk == 1 else None
        if lone_side is None:
            new = None
        else:
            other_m, other_k = (bm, bk) if lone_side == 0 else (wm, wk)
            lone_mask = b.white_kings if lone_side == 0 else b.black_kings
            diagonal = any((lone_mask >> s) & 1 for s in range(32) if NAMES[s][0] == str("abcdefgh"[int(NAMES[s][1]) - 1]))
            if other_k >= 1 and other_m + other_k <= 2:
                limit = 5
            elif other_k >= 1 and other_m + other_k == 3 and diagonal:
                limit = 5
            elif other_m == 0 and other_k >= 3:
                limit = 15
            else:
                limit = None
            new = (lone_side, other_m, other_k, limit) if limit is not None else None
        if new != self.regulation_key:
            self.regulation_key = new
            self.regulation_moves = [0, 0]

    def update(self, board, move, moving_piece, mover):
        self.plies += 1
        self.seen[key(board)] += 1
        self.quiet_kings = self.quiet_kings + 1 if abs(moving_piece) == 2 and not move.captured_list else 0
        self.material_stable = 0 if move.captured_list or move.is_promotion else self.material_stable + 1
        old = self.regulation_key
        self._regulation(board)
        if old is not None and old == self.regulation_key:
            self.regulation_moves[0 if mover == Color.WHITE else 1] += 1

    def outcome(self, board, legal):
        if not legal:
            return -board.turn.value, "no_legal_moves"
        if self.seen[key(board)] >= 3:
            return 0, "threefold"
        if self.quiet_kings >= 30:
            return 0, "15_king_moves"
        if self.regulation_key is not None:
            lone, _, _, limit = self.regulation_key
            if self.regulation_moves[1 - lone] >= limit:
                return 0, f"regulation_{limit}"
        count = (board.white_men | board.white_kings | board.black_men | board.black_kings).bit_count()
        both_kings = bool(board.white_kings and board.black_kings)
        if both_kings and ((count <= 5 and self.material_stable >= 60) or (count <= 7 and self.material_stable >= 120)):
            return 0, "stable_endgame"
        if self.plies >= self.max_plies:
            return 0, "experimental_ply_cap"
        return None