__author__ = 'Tom Schaul, tom@idsia.ch'

from pybrain.rl.environments.twoplayergames.gomoku import GomokuGame


class PenteGame(GomokuGame):
    """ The game of Pente.

    The rules are similar to Go-Moku, except that it is now possible to capture
    stones, in pairs, by putting stones at both ends of a pair of the opponent.
    The game is won by the first player who either has 5 connected stones, or
    has captured 5 pairs.
    """

    def reset(self):
        GomokuGame.reset(self)
        self.pairsTaken = {self.BLACK: 0, self.WHITE: 0}
        center = (self.size[0] // 2, self.size[1] // 2)
        self._setStone(-self.startcolor, center)
        self.movesDone += 1

    def getKilling(self, c):
        """ return all legal positions for a color that immediately kill the opponent. """
        res = GomokuGame.getKilling(self, c)
        for p in self.getLegals(c):
            k = self._killsWhich(c, p)
            if self.pairsTaken[c] + len(k) / 2 >= 5:
                res.append(p)
        return res

    def _killsWhich(self, c, pos):
        """ placing a stone of color c at pos would kill which enemy stones? """
        res = []
        for dir in [(0, 1), (1, 0), (1, 1), (1, -1)]:
            for d in [-1, 1]:
                killcands = []
                for i in [1, 2, 3]:
                    next = (pos[0] + dir[0] * i * d, pos[1] + dir[1] * i * d)
                    if (next[0] < 0 or next[0] >= self.size[0]
                        or next[1] < 0 or next[1] >= self.size[1]):
                        break
                    if i == 3 and self.b[next] == c:
                        res += killcands
                        break
                    if i != 3 and self.b[next] != -c:
                        break
                    killcands.append(next)
        return res

    def doMove(self, c, pos):
        """ the action is a (color, position) tuple, for the next stone to move.
        returns True if the move was legal. """
        self.movesDone += 1
        if not self.isLegal(c, pos):
            return False
        elif self._fiveRow(c, pos):
            self.winner = c
            self.b[pos] = 'x'
            return True
        else:
            tokill = self._killsWhich(c, pos)
            if self.pairsTaken[c] + len(tokill) / 2 >= 5:
                self.winner = c
                self.b[pos] = 'x'
                return True

            self._setStone(c, pos, tokill)
            if self.movesDone == (self.size[0] * self.size[1]
                                  + 2 * (self.pairsTaken[self.BLACK] + self.pairsTaken[self.WHITE])):
                # DRAW
                self.winner = self.DRAW
            return True

    def _setStone(self, c, pos, tokill=None):
        """ set stone, and potentially kill stones. """
        if tokill == None:
            tokill = self._killsWhich(c, pos)
        GomokuGame._setStone(self, c, pos)
        for p in tokill:
            self.b[p] = self.EMPTY
        self.pairsTaken[c] += len(tokill) // 2

    def __str__(self):
        s = GomokuGame.__str__(self)
        s += 'Black captured:' + str(self.pairsTaken[self.BLACK]) + ', white captured:' + str(self.pairsTaken[self.WHITE]) + '.'
        return s

