pygo/pygo.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# A GTK Python GO client

import os
import pygame
from pygame.locals import *


def load_png(name):
    """ Load image and return image object"""
    fullname = os.path.join('res', name)

    try:
        image = pygame.image.load(fullname)
        if image.get_alpha() is None:
            image = image.convert()
        else:
            image = image.convert_alpha()

    except pygame.error, message:
        print 'Cannot load image:', fullname
        raise SystemExit, message

    return image


def build_img_res():
    ret = {}

    triangle = load_png('go_t.png')
    
    ret['w'] = load_png('go_w.png')
    ret['wT'] = load_png('go_w.png')
    ret['wT'].blit(triangle, (0,0))

    ret['b'] = load_png('go_b.png')
    ret['bT'] = load_png('go_b.png')
    ret['bT'].blit(triangle, (0,0))

    circle = load_png('go_c.png')

    for d in ('u', 'd', 'l', 'r', 'm', 'dl', 'dr', 'ul', 'ur', 'h'):
        ret[d] = load_png('go_' + d + '.png')
        ret[d + 'C'] = load_png('go_' + d + '.png')
        ret[d + 'C'].blit(circle, (0,0))

    return ret


class GobanSquare:
    """A single square on the go board"""

    def __init__(self, pos):
        self.x, self.y = pos
        self.state = -1
        self.marked = False
        self.checked = False # Used for recursive checks

        if (self.x, self.y) == (1,1):
            self.default_draw_code = 'ul'
        elif (self.x, self.y) == (1,19):
            self.default_draw_code = 'ur'
        elif (self.x, self.y) == (19,1):
            self.default_draw_code = 'dl'
        elif (self.x, self.y) == (19,19):
            self.default_draw_code = 'dr'
        elif (self.x, self.y) in [(4,4), (4,10), (4,16), (10,4), (10,10), (10,16), (16,4), (16,10), (16,16)]:
            self.default_draw_code = 'h'
        elif self.x == 1:
            self.default_draw_code = 'u'
        elif self.y == 1:
            self.default_draw_code = 'l'
        elif self.x == 19:
            self.default_draw_code = 'd'
        elif self.y == 19:
            self.default_draw_code = 'r'
        else:
            self.default_draw_code = 'm'


    def toggle_marked(self):
        if self.marked:
            self.marked = False
        else:
            self.marked = True


    def get_draw_code(self):
        ret = None

        if self.state == -1:
            ret = self.default_draw_code
        elif self.state == 0:
            ret = 'b'
        elif self.state == 1:
            ret = 'w'

        if self.marked:
            if self.state >=0:
                ret = self.default_draw_code + 'T'
            else:
                ret = self.default_draw_code + 'C'

        return ret



class Goban:
    """Represents the go board. Handles stone placement, captures, etc"""

    def __init__(self):
        # Build the 361 board intersections
        self.board = []
        for i in range(19):
            self.board.append([])
            for j in range(19):
                self.board[i].append(GobanSquare((i+1, j+1)))

        self.turn = 0
        self.captures = []
        self.captures.append(0)
        self.captures.append(0)


    def place_stone(self, pos):
        if not self._valid_move(pos):
            return

        x, y = pos
        self.board[x][y].state = self.turn
        self._capture(pos)
        self.turn = (self.turn + 1) % 2


    def toggle_marked(self, pos):
        x,y = pos
        if x < 0 or x > 18 or y < 0 or y > 18:
            return
        
        self.board[x][y].toggle_marked()
        

    def _capture(self, pos):
        x, y = pos
        who = (self.turn + 1) % 2

        for p in [(x, y+1), (x, y-1), (x+1, y), (x-1, y)]:
            newx, newy = p
            if newx < 0 or newx > 18 or newy < 0 or newy > 18 or self.board[newx][newy].state == self.turn:
                continue
                
            if not self._has_liberties(p, who):
                self.captures[self.turn] += self._delete_group(p)

            self._clear_checks()


    def _valid_move(self, pos):
        x, y = pos
        if x < 0 or x > 18 or y < 0 or y > 18:
            return False

        # Can't play atop another stone
        if self.board[x][y].state != -1:
            return False

        # Temporarily place the stone
        self.board[x][y].state = self.turn

        liberties = self._has_liberties(pos, self.turn)
        self._clear_checks()

        opponent = (self.turn + 1) % 2

        kills_group = False
        for d in [(x, y+1), (x, y-1), (x+1, y), (x-1, y)]:
            if not self._has_liberties(d, opponent):
                kills_group = True
                break
            self._clear_checks()

        # Remove temporary stone
        self.board[x][y].state = -1

        return liberties > 0 or kills_group


    # Recursively find whether there are liberties for the group
    # at pos. Positive numbers are not necessarily accurate - 
    # treat this as a boolean
    def _has_liberties(self, pos, who, direction = None):
        x,y = pos
        if x < 0 or x > 18 or y < 0 or y > 18:
            return 0

        square = self.board[x][y]

        if square.checked:
            return 0
        else:
            square.checked = True

        if square.state == -1:
            return 1
        elif square.state != who:
            return 0
        else:
            liberties = 0
            for d in [(x, y+1), (x-1, y), (x+1, y), (x, y-1)]:
                liberties += self._has_liberties(d, who)

            return liberties


    # We don't need to worry about crossing ourselves with the
    # recursion here, because we've already deleted the stone.
    # It would be more efficient to avoid going backwards,
    # but a lot more complicated (see _has_liberties)
    def _delete_group(self, pos):
        x,y = pos
        if x < 0 or x > 18 or y < 0 or y > 18:
            return

        who = self.board[x][y].state

        if who == -1:
            return

        self.board[x][y].state = -1

        num_deleted = 1

        num_deleted += self._delete_group_r((x, y+1), who)
        num_deleted += self._delete_group_r((x, y-1), who)
        num_deleted += self._delete_group_r((x+1, y), who)
        num_deleted += self._delete_group_r((x-1, y), who)

        return num_deleted


    def _delete_group_r(self, pos, who):
        x,y = pos
        if x < 0 or x > 18 or y < 0 or y > 18:
            return 0

        if self.board[x][y].state != who:
            return 0

        self.board[x][y].state = -1

        num_deleted = 1

        num_deleted += self._delete_group_r((x, y+1), who)
        num_deleted += self._delete_group_r((x, y-1), who)
        num_deleted += self._delete_group_r((x+1, y), who)
        num_deleted += self._delete_group_r((x-1, y), who)

        return num_deleted


    def _clear_checks(self):
        for row in self.board:
            for col in row:
                col.checked = False


    def draw_board(self, size, img_res):
        ret = pygame.Surface((size,size))

        inc = size / 19;
        i = 0
        for row in self.board:
            j = 0
            for square in row:
                s = pygame.transform.scale(img_res[square.get_draw_code()], (inc, inc))
                ret.blit(s, (j*inc,i*inc))
                j += 1
            i += 1

        return ret.convert()


    def draw_info(self):
        textbox = pygame.Surface((150, 300))
        textbox = textbox.convert()
        textbox.fill((250, 250, 250))

        font = pygame.font.Font(None, 24)
        heading = font.render('Captures', 1, (10, 10, 10))
        black_cap = font.render('Black: {}'.format(self.captures[0]), 1, (10, 10, 10))
        white_cap = font.render('White: {}'.format(self.captures[1]), 1, (10, 10, 10))
        turn = font.render('Turn: {}'.format(['Black', 'White'][self.turn]), 1, (10, 10, 10))

        textbox.blit(heading, (0, 0))
        textbox.blit(black_cap, (0, 28))
        textbox.blit(white_cap, (0, 56))
        textbox.blit(turn, (0, 100))

        return textbox



def main():
    # Basic screen init
    pygame.init()
    screen = pygame.display.set_mode((1000, 800))
    pygame.display.set_caption('pyGo')
    
    # Create the background object, make it blank
    background = pygame.Surface(screen.get_size())
    background = background.convert()
    background.fill((250, 250, 250))

    # Build the dict of image objects
    img_res = build_img_res()
            
    board_size = 800
    board_inc = board_size / 19

    goban = Goban()

    board = goban.draw_board(board_size, img_res)
    background.blit(board, (0,0))

    text = goban.draw_info()
    background.blit(text, (815, 25))
    
    screen.blit(background, (0, 0))
    pygame.display.flip()

    while True:
        event = pygame.event.wait()
        
        if event.type == QUIT:
            return
        if event.type == MOUSEBUTTONDOWN:
            x, y = event.pos
            row = y / board_inc
            col = x / board_inc

            if x <= board_size:
                if event.button == 1:
                    goban.place_stone((row, col))
                if event.button == 3:
                    goban.toggle_marked((row, col))
                
        board = goban.draw_board(board_size, img_res)
        background.blit(board, (0,0))
    
        text = goban.draw_info()
        background.blit(text, (815, 25))

        screen.blit(background, (0, 0))
        pygame.display.flip()

    

if __name__ == '__main__': main()