chess-python/tree.py at main · pacanada/chess-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import json
import time
from copy import deepcopy
from typing import Optional

from chess_python.chess import Chess

DICT_POSITIONS = {
    0: {
        "name": "init_position",
        "depth": 3,
        "fen": "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1",
    },
    # 1: {
    #     "name": "pos_3",
    #     "depth": 3,
    #     "fen": "8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - - 0 1",
    # },
    # 2: {
    #     "name": "pos_4",
    #     "depth": 3,
    #     "fen": "r3k2r/Pppp1ppp/1b3nbN/nP6/BBP1P3/q4N2/Pp1P2PP/R2Q1RK1 w kq - 0 1",
    # },
    # 3: {
    #     "name": "pos_5",
    #     "depth": 3,
    #     "fen": "rnbq1k1r/pp1Pbppp/2p5/8/2B5/8/PPP1NnPP/RNBQK2R w KQ - 1 8",
    # },
    # 4: {
    #     "name": "pos_3",
    #     "depth": 4,
    #     "fen": "8/2p5/3p4/KP5r/1R3p1k/8/4P1P1/8 w - - 0 1",
    # },
    # 5: {
    #     "name": "init_position",
    #     "depth": 4,
    #     "fen": "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1",
    # },
    6: {
        "name": "puzzle",
        "depth": 3,
        "fen": "3qr2k/pbpp2pp/1p5N/3Q2b1/2P1P3/P7/1PP2PPP/R4RK1 w - - 0 1",
    },
    7: {"name": "bug", "depth": 1, "fen": "2N3k1/1P6/1B2P1Bp/7p/8/2b1r3/4R3/4K3 w - - 0 1"},
}


class Node:
    def __init__(self, game: Chess, depth: int, move: Optional[str] = None):
        self.move = move
        # no need to copy the entire game for the last branch since we are not doing any move
        game = deepcopy(game) if depth != 0 else game
        self.children = get_children(game, depth, move)


def get_children(game, depth, move):
    if depth == 0:
        return []
    if move is not None:
        game.move(move)
    # for ordering and comparing
    # allowed_moves_i_f= sorted(game.legal_moves())
    allowed_moves_i_f = game.legal_moves()
    children = []
    for move in allowed_moves_i_f:
        children.append(Node(game, depth - 1, move))
    return children


def get_size_node(node):
    if len(node.children) == 0:
        return 1
    else:
        return sum([get_size_node(child) for child in node.children])


def get_move_combination(node):
    if len(node.children) == 0:
        return {}
    else:
        return {child.move: get_move_combination(child) for child in node.children}


if __name__ == "__main__":
    for _, v in DICT_POSITIONS.items():
        print("Evaluating position {}".format(v["name"]))
        t1 = time.time()
        node = Node(Chess(v["fen"]), v["depth"])  # type: ignore[arg-type]
        t2 = time.time()
        print(f"Size of own tree {get_size_node(node)}, Elapsed time {(t2-t1):.2f} s")
        move_combination = get_move_combination(node)
        with open("moves_own.json", "w") as f:
            json.dump(move_combination, f)