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

def _alphabeta(
    board: Goban.Board,
    heuristic: Callable[[Goban.Board, Any], float],
    color,
    alpha=-math.inf,
    beta=math.inf,
    depth: int = 3,
    shouldStop=lambda: False
) -> tuple[float, Any]:
    if board.is_game_over() or depth == 0:
        return heuristic(board, color), None

    wantMax = (board.next_player() == color)
    best_move = None

    if wantMax:
        acc = -math.inf
        for move in board.generate_legal_moves():
            if Goban.Board.flat_to_name(move) == "PASS":
                continue

            board.push(move)
            value = _alphabeta(
                board,
                heuristic=heuristic,
                color=color,
                alpha=alpha,
                beta=beta,
                depth=depth - 1,
                shouldStop=shouldStop
            )[0]
            board.pop()

            if value > acc:
                acc = value
                best_move = move

            alpha = max(alpha, acc)
            if shouldStop() or acc >= beta:
                break  # beta cutoff

    else:
        acc = math.inf
        for move in board.generate_legal_moves():
            if Goban.Board.flat_to_name(move) == "PASS":
                continue

            board.push(move)
            value = _alphabeta(
                board,
                heuristic=heuristic,
                color=color,
                alpha=alpha,
                beta=beta,
                depth=depth - 1,
                shouldStop=shouldStop
            )[0]
            board.pop()

            if value < acc:
                acc = value
                best_move = move

            beta = min(beta, acc)
            if shouldStop() or acc <= alpha:
                break  # alpha cutoff

    return acc, best_move

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


def IDDFS(
    board: Goban.Board,
    heuristic: Callable[[Goban.Board, Any], float],
    color,
    max_depth: int = 10,
    duration: float = 5.0  # Duration in seconds
):
    best_move = None
    start_time = time.time()
    shouldStop = lambda: (time.time() - start_time) >= duration

    for depth in range(1, max_depth + 1):

        value, move = _alphabeta(
            board,
            heuristic=heuristic,
            color=color,
            depth=depth,
            shouldStop=shouldStop
        )

        if shouldStop():
            break

        print(f"{depth}, {value}", file=stderr)
        best_move = move

    return best_move