1import collections
2
3from .ChessFile import ChessFile, LoadingError
4from pychess.Utils.GameModel import GameModel
5from pychess.Utils.const import WHITE, BLACK, WON_RESIGN, WAITING_TO_START, BLACKWON, WHITEWON, DRAW, FISCHERRANDOMCHESS
6from pychess.Utils.logic import getStatus
7from pychess.Utils.lutils.leval import evaluateComplete
8from pychess.Variants.fischerandom import FischerandomBoard
9
10__label__ = _("Chess Position")
11__ending__ = "epd"
12__append__ = True
13
14
15def save(handle, model, position=None, flip=False):
16    """Saves game to file in fen format"""
17
18    color = model.boards[-1].color
19
20    fen = model.boards[-1].asFen().split(" ")
21
22    # First four parts of fen are the same in epd
23    handle.write(" ".join(fen[:4]))
24
25    ############################################################################
26    # Repetition count                                                         #
27    ############################################################################
28    rep_count = model.boards[-1].board.repetitionCount()
29
30    ############################################################################
31    # Centipawn evaluation                                                     #
32    ############################################################################
33    if model.status == WHITEWON:
34        if color == WHITE:
35            ceval = 32766
36        else:
37            ceval = -32766
38    elif model.status == BLACKWON:
39        if color == WHITE:
40            ceval = -32766
41        else:
42            ceval = 32766
43    elif model.status == DRAW:
44        ceval = 0
45    else:
46        ceval = evaluateComplete(model.boards[-1].board, model.boards[-1].color)
47
48    ############################################################################
49    # Opcodes                                                                  #
50    ############################################################################
51    opcodes = (
52        ("fmvn", fen[5]),  # In fen full move number is the 6th field
53        ("hmvc", fen[4]),  # In fen halfmove clock is the 5th field
54
55        # Email and name of receiver and sender. We don't know the email.
56        ("tcri", "?@?.? %s" % repr(model.players[color]).replace(";", "")),
57        ("tcsi", "?@?.? %s" % repr(model.players[1 - color]).replace(";", "")),
58        ("ce", ceval),
59        ("rc", rep_count), )
60
61    for key, value in opcodes:
62        handle.write(" %s %s;" % (key, value))
63
64    ############################################################################
65    # Resign opcode                                                            #
66    ############################################################################
67    if model.status in (WHITEWON, BLACKWON) and model.reason == WON_RESIGN:
68        handle.write(" resign;")
69
70    print("", file=handle)
71    handle.close()
72
73
74def load(handle):
75    return EpdFile(handle)
76
77
78class EpdFile(ChessFile):
79    def __init__(self, handle):
80        ChessFile.__init__(self, handle)
81
82        self.games = [self.create_rec(line.strip()) for line in handle if line]
83        self.count = len(self.games)
84
85    def create_rec(self, line):
86        rec = collections.defaultdict(str)
87        rec["Id"] = 0
88        rec["Offset"] = 0
89        rec["FEN"] = line
90
91        castling = rec["FEN"].split()[2]
92        for letter in castling:
93            if letter.upper() in "ABCDEFGH":
94                rec["Variant"] = FISCHERRANDOMCHESS
95                break
96
97        return rec
98
99    def loadToModel(self, rec, position, model=None):
100        if not model:
101            model = GameModel()
102
103        if "Variant" in rec:
104            model.variant = FischerandomBoard
105
106        fieldlist = rec["FEN"].split(" ")
107        if len(fieldlist) == 4:
108            fen = rec["FEN"]
109            opcodestr = ""
110
111        elif len(fieldlist) > 4:
112            fen = " ".join(fieldlist[:4])
113            opcodestr = " ".join(fieldlist[4:])
114
115        else:
116            raise LoadingError("EPD string can not have less than 4 field")
117
118        opcodes = {}
119        for opcode in map(str.strip, opcodestr.split(";")):
120            space = opcode.find(" ")
121            if space == -1:
122                opcodes[opcode] = True
123            else:
124                opcodes[opcode[:space]] = opcode[space + 1:]
125
126        if "hmvc" in opcodes:
127            fen += " " + opcodes["hmvc"]
128        else:
129            fen += " 0"
130
131        if "fmvn" in opcodes:
132            fen += " " + opcodes["fmvn"]
133        else:
134            fen += " 1"
135
136        model.boards = [model.variant(setup=fen)]
137        model.variations = [model.boards]
138        model.status = WAITING_TO_START
139
140        # rc is kinda broken
141        # if "rc" in opcodes:
142        #    model.boards[0].board.rc = int(opcodes["rc"])
143
144        if "resign" in opcodes:
145            if fieldlist[1] == "w":
146                model.status = BLACKWON
147            else:
148                model.status = WHITEWON
149            model.reason = WON_RESIGN
150
151        if model.status == WAITING_TO_START:
152            status, reason = getStatus(model.boards[-1])
153            if status in (BLACKWON, WHITEWON, DRAW):
154                model.status, model.reason = status, reason
155
156        return model
157
158    def get_player_names(self, rec):
159        data = rec["FEN"]
160
161        names = {}
162
163        for key in "tcri", "tcsi":
164            keyindex = data.find(key)
165            if keyindex == -1:
166                names[key] = _("Unknown")
167            else:
168                sem = data.find(";", keyindex)
169                if sem == -1:
170                    opcode = data[keyindex + len(key) + 1:]
171                else:
172                    opcode = data[keyindex + len(key) + 1:sem]
173                name = opcode.split(" ", 1)[1]
174                names[key] = name
175
176        color = data.split(" ")[1] == "b" and BLACK or WHITE
177
178        if color == WHITE:
179            return (names["tcri"], names["tcsi"])
180        else:
181            return (names["tcsi"], names["tcri"])
182