from sys import stderr
import time
import math
from typing import Any, Callable
import Goban


def _next_color(color):
    return Goban.Board._BLACK if color == Goban.Board._WHITE else Goban.Board._WHITE

# Returns heuristic, move
def _alphabeta(
    board: Goban.Board,
    heuristic: Callable[[Goban.Board, Any], float],
    color,
    move,
    alpha=-math.inf,
    beta=math.inf,
    depth: int = 3,
) -> tuple[float, Any]:

    wantMax = (board.next_player != color)
    if depth == 0 or board.is_game_over():
        return heuristic(board, color), move

    if wantMax:
        acc = -math.inf, None
        for move in board.generate_legal_moves():
            board.push(move)
            value = (
                _alphabeta(
                    board,
                    alpha=alpha,
                    beta=beta,
                    move=move,
                    heuristic=heuristic,
                    color=_next_color(color),
                    depth=depth - 1,
                )[0],
                move,
            )
            acc = max(
                acc,
                value,
                key=lambda t: t[0],
            )
            board.pop()

            if acc[0] >= beta:
                break  # beta cutoff
            alpha = max(alpha, value[0])

    else:
        acc = math.inf, None
        for move in board.generate_legal_moves():
            board.push(move)
            value = (
                _alphabeta(
                    board,
                    alpha=alpha,
                    beta=beta,
                    move=move,
                    heuristic=heuristic,
                    color=_next_color(color),
                    depth=depth - 1,
                )[0],
                move,
            )
            acc = min(
                acc,
                value,
                key=lambda t: t[0],
            )
            board.pop()

            if acc[0] <= alpha:
                break  # alpha cutoff
            beta = min(beta, value[0])

    return acc


def alphabeta(
    board: Goban.Board,
    heuristic: Callable[[Goban.Board, Any], float],
    color,
    depth: int = 3,
):
    _, move = _alphabeta(board, move=-1, heuristic=heuristic, color=color, depth=depth)
    return move


def IDDFS(board: Goban.Board, heuristic, color, duration: float, maxdepth=42):
    st = time.time()
    depth = 1
    move = -1

    while time.time() - st < duration and depth <= maxdepth:
        print("depth:", depth, time.time() - st, file=stderr)
        move = _alphabeta(
            board, heuristic, color, move=-1, alpha=-10, beta=10, depth=depth
        )[1]
        depth += 1

    print(time.time() - st, duration, depth, file=stderr)
    return move
