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,
    shouldStop = lambda: False
) -> 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():
            if Goban.Board.flat_to_name(move) == "PASS":
                continue

            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 shouldStop() or acc[0] >= beta:
                break  # beta cutoff
            alpha = max(alpha, acc[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 shouldStop() or acc[0] <= alpha:
                break  # alpha cutoff
            beta = min(beta, acc[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()
    shouldStop = (lambda: time.time() - st > duration)
    depth = 0
    move = -1
    score = -1

    while not shouldStop() and depth <= maxdepth:
        if depth % 2 == 0:
            score, move = _alphabeta(
                    board, heuristic, color, move=move, alpha=-math.inf, beta=math.inf, depth=depth, shouldStop=shouldStop
            )

            if score == math.inf:
                return move, score

        else:
            score, move = _alphabeta(
                board, heuristic, color, move=move, alpha=-math.inf, beta=math.inf, depth=depth, shouldStop=shouldStop
            )

            if score == -math.inf:
                return move, score

        print("depth:", depth, time.time() - st, score,  file=stderr)
        depth += 1

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