# coding: utf-8
"""
A simple connect four game with AI using alpha-beta pruning

(c) Julian Habrock <jug@fantasymail.de>, 9/2009
    bytemuehle.de
"""

import pygame
import pygame.locals as loc

#########
# ai part

class Grid(object):
    width = 7
    height = 6
    connect = 4
    def __init__(self, grid=None):
        if grid:
            if type(grid) == list and len(grid) == self.width * self.height:
                self._data = list
            elif type(grid) == Grid:
                self._data = grid._data[:]
                self.width, self.height = grid.width, grid.height
            else:
                raise ValueError, ("grid argument must be either a list" + 
                    "of length %s or another grid" % self.width * self.height)
        else:
            self._data = [0] * self.width * self.height

    def drop(self, col, value):
        if not 0 <= col < self.width:
            raise ValueError, "%i must be between 0 and %i" % (col, self.width)
        last_free_position = None
        for i in xrange(self.height):
            if not self._data[i * self.width + col]:
                last_free_position = i
            else:
                break
        if last_free_position is None:
            return False
        self._data[last_free_position * self.width + col] = value
        return True
    
    def copy(self):
        return Grid(self)

    def score(self):
        # number possible win positions for each player with
        # using at least one of the current chips
        possible_positions = [0, 0]
        for line in self:
            if len(line) < self.connect:
                continue
            for i in xrange(len(line)-self.connect+1):
                slice = line[i:i+self.connect]
                if abs(sum(slice)) == self.connect:
                    return 100*slice[0]
                # count only possible positions with at least one chip already
                # set. Completely free positions are the same for both players.
                #! double score positions that are nearly finished?
                if set(slice) == set([0, 1]):
                    possible_positions[0] += 1
                elif set(slice) == set([0, -1]):
                    possible_positions[0] += 1
        return possible_positions[0] - possible_positions[1]

    def is_full(self):
        return not 0 in self._data

    def iter_rows(self):
        for i in xrange(self.height):
            yield self._data[i*self.width:(i+1)*self.width]

    def iter_cols(self):
        for i in xrange(self.width):
            yield self._data[i::self.width]

    def iter_lr_diagonals(self):
        for i in xrange(self.height-1, 0, -1):
            v = self._data[i*self.width:
                             i*self.width+(self.height-i)*(self.width+1):
                             self.width+1]
            yield v if len(v) <= self.width else v[:self.width]
        for i in xrange(self.width):
            v = self._data[i:(self.width-i)*(self.width+1):self.width+1]
            yield v if len(v) <= self.height else v[:self.height]

    def iter_rl_diagonals(self):
        for i in xrange(self.width):
            v = self._data[i:i*(self.width)+1:self.width-1]
            yield v if len(v) <= self.height else v[:self.height]
        for i in xrange(2, self.height+1):
            v = self._data[i*self.width-1:
                             i*self.width+(self.height-i)*(self.width-1):
                             self.width-1]
            yield v if len(v) <= self.width else v[:self.width]

    def __iter__(self):
        for iterator in (self.iter_cols, self.iter_rows,
                         self.iter_lr_diagonals, self.iter_rl_diagonals):
            for line in iterator():
                yield line


    def __str__(self):
        out = ""
        chars = {0:"_", 1:"0", -1:"X"}
        for row in self.iter_rows():
            for c in row:
                out += chars[c] if c in chars else str(c)
            out += "\n"
        return out

    def raw(self):
        print "raw data"
        print "\n".join([" ".join([(c and "%2.i" % c or " 0") for c in line])
                                                  for line in self.iter_rows()])
#        print
#        print "\n".join([" ".join([(c and "%2.i" % c or " 0") for c in line])
#                                                  for line in self.iter_cols()])
#        print
#        print "left to right diagonals"
#        print "\n".join([" ".join([(c and "%2.i" % c or " 0") for c in line])
#                                                  for line in self.iter_lr_diagonals()])
#        print
#        print "right to left diagonals"
#        print "\n".join([" ".join([(c and "%2.i" % c or " 0") for c in line])
#                                                  for line in self.iter_rl_diagonals()])

# test it

#g = Grid()
#print type(g)
#print g
#g.drop(3, 1)
#print g
#g.drop(4, -1)
#print g
#g.drop(3, 1)
#print g
#g.drop(5, 1)
#print g

#print len([x for x in g])
#print
#g.raw()
#for a, b in ((2,6),(5,5), (6,2)):
#    print "-"*20
#    print a, b
#    g.width = a
#    g.height = b
#    g.raw()




def max(deep, alpha, beta, grid, first=True):
    """
    The max-function for the minimax algorithm.
    """
    score = grid.score()
    if deep == 0 or abs(score) == 100 or grid.is_full():
        # give earlier wins a higher score
        if score == 100:
            return 100 + deep
        elif score == -100:
            return -100 - deep
        return score
    best_i = None
    for i in xrange(grid.width):
        temp_grid = grid.copy()
        if temp_grid.drop(i, 1):
            s = min(deep-1, alpha, beta, temp_grid, False)
            if s >= beta:
                if first:
                    return i
                else:
                    return beta
            if s > alpha:
                alpha = s
                best_i = i
    if first:
        return best_i
    else:
        return alpha


def min(deep, alpha, beta, grid, first=True):
    """
    The min-function for the minimax algorithm.
    Since this is nearly the same as ``max`` both methods should be
    replaced with one negamax function.
    """
    score = grid.score()
    if deep == 0 or abs(score) == -100 or grid.is_full():
        if score == -100:
            return -100 - deep
        elif score == 100:
            return 100 + deep
        return score
    best_i = None
    for i in xrange(grid.width):
        temp_grid = grid.copy()
        if temp_grid.drop(i, -1):
            s = max(deep-1, alpha, beta, temp_grid, False)
            if s <= alpha:
                if first:
                    return i
                else:
                    return alpha
            if s < beta:
                beta = s
                best_i = i
    if first:
        return best_i
    else:
        return beta


def ai(grid):
    """
    wrap the max function call to make it even easier to call it
    from the game part
    """
    return max(4, -1000, 1000, grid)


##########
# gamepart


class Game(object):
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((675, 550))
        pygame.display.set_caption('(0nn3(tF0UR - by jug')
            
        self.chip_colors = {1: (200, 0, 0), -1: (0, 200, 0)}

        if pygame.font:
            font = pygame.font.Font(None, 36)
            text = font.render("ConnectF0UR", 1, (255,255,255), (0,0,0))
            textpos = text.get_rect(centerx=self.screen.get_width()/2, top=20)
            self.screen.blit(text, textpos)
            font = pygame.font.Font(None, 20)
            font.set_italic(True)
            text = font.render("press space to restart", 1, (255,255,255), (0,0,0))
            textpos = text.get_rect(centerx=self.screen.get_width()/2, top=450)
            self.screen.blit(text, textpos)
            
        self.clock = pygame.time.Clock()
        
        self.grid = Grid()
        self.status = ""
        stati = {0:"draw", -10:"you win", 10:"you lose"}
        self.draw_grid_background()
        self.run()
        
    def draw_grid_background(self):
        self.screen.fill((0,0,0), (75,75,525,450))
        fg = (255,255,255)
        for i in xrange(75, 601, 75):
            pygame.draw.line(self.screen, fg, (i, 65), (i, 525)) 
            if i <= 525:
                pygame.draw.line(self.screen, fg, (75, i), (600, i)) 
        pygame.display.flip()
        
    def run(self):
        while 1:
            drects = []
            self.clock.tick(100)
            for event in pygame.event.get():
                if event.type == loc.QUIT:
                    return
                elif event.type == loc.KEYDOWN:
                    if event.key == loc.K_ESCAPE:
                        return
                    elif event.key == loc.K_SPACE:
                        self.draw_grid_background()
                        self.grid = Grid()
                        status = ""
                    elif event.key == loc.K_s:
                        print self.grid.score()
                if event.type == loc.MOUSEBUTTONDOWN and not self.status:
                    x, y = event.pos
                    if 75 < x < 600 and 75 < y < 525:
                        index = (x-75)/75
                        if self.grid.drop(index, -1):
                            ai_chip = ai(self.grid)
                            if ai_chip > 100:
                                print "you lose"
                            elif ai_chip < -100:
                                print "you win"
                            else:
                                self.grid.drop(ai_chip, 1)
            for y, row in enumerate(self.grid.iter_rows()):
                for x, value in enumerate(row):
                    if value:
                        drects.append(pygame.draw.circle(self.screen,
                                                 self.chip_colors[value],
                                        (int(75*(x+1.5)), int(75*(y+1.5))), 30))
            pygame.display.update(drects)
            

if __name__ == '__main__':
    Game()