# Graphical Four-In-A-Row (a Connect Four clone)
# http://inventwithpython.com/blog
# By Al Sweigart al@inventwithpython.com

import random
import copy
import sys
import pygame
from pygame.locals import *

BOARDWIDTH = 7  # how many spaces wide the board is
BOARDHEIGHT = 6 # how many spaces tall the board is

DIFFICULTY = 2 # how many moves to look ahead. > 2 usually takes too long to compute since it's exponential

SPACESIZE = 50 # The size of the tokens and individual board spaces

FPS = 30 # frames per second to update the screen
WINDOWWIDTH = 640 # width of the program's window, in pixels
WINDOWHEIGHT = 480 # height in pixels

XMARGIN = int((WINDOWWIDTH - BOARDWIDTH * SPACESIZE) / 2) # approx. number of pixels on the left or right side of the board to the edge of the window
YMARGIN = int((WINDOWHEIGHT - BOARDHEIGHT * SPACESIZE) / 2) # approx. number of pixels above or below board

BRIGHTBLUE = (0, 50, 255)
WHITE = (255, 255, 255)

BGCOLOR = BRIGHTBLUE
TEXTCOLOR = WHITE

redPileRect = pygame.Rect(int(SPACESIZE / 2), WINDOWHEIGHT - int(3 * SPACESIZE / 2), SPACESIZE, SPACESIZE)
blackPileRect = pygame.Rect(WINDOWWIDTH - int(3 * SPACESIZE / 2), WINDOWHEIGHT - int(3 * SPACESIZE / 2), SPACESIZE, SPACESIZE)

def main():
    global redTokenImg, blackTokenImg, boardImg, gameClock, windowSurf

    pygame.init()
    gameClock = pygame.time.Clock()
    windowSurf = pygame.display.set_mode((WINDOWWIDTH, WINDOWHEIGHT))
    pygame.display.set_caption('Four in a Row')

    redTokenImg = pygame.image.load('4row_red.png')
    redTokenImg = pygame.transform.smoothscale(redTokenImg, (SPACESIZE, SPACESIZE))
    blackTokenImg = pygame.image.load('4row_black.png')
    blackTokenImg = pygame.transform.smoothscale(blackTokenImg, (SPACESIZE, SPACESIZE))
    boardImg = pygame.image.load('4row_board.png')
    boardImg = pygame.transform.smoothscale(boardImg, (SPACESIZE, SPACESIZE))

    bigFont = pygame.font.Font('freesansbold.ttf', 72)
    normalFont = pygame.font.Font('freesansbold.ttf', 24)

    gameOverSurf2 = normalFont.render('Click to start a new game.', 1, TEXTCOLOR)
    gameOverRect2 = gameOverSurf2.get_rect()
    gameOverRect2.center = (int(WINDOWWIDTH / 2), int(WINDOWHEIGHT / 2) + 50)

    isFirstGame = True

    while True:
        if isFirstGame:
            turn = 'computer'
            isFirstGame = False
        else:
            if random.randint(0, 1) == 0:
                turn = 'computer'
            else:
                turn = 'human'
        mainBoard = getNewBoard()

        while True:
            if turn == 'human':
                getHumanMove(mainBoard)
                if isWinner(mainBoard, 'red'):
                    winner = 'Human Wins'
                    break
                turn = 'computer'
            else:
                column = getComputerMove(mainBoard)
                animateComputerMoving(mainBoard, column)
                makeMove(mainBoard, 'black', column)
                if isWinner(mainBoard, 'black'):
                    winner = 'Computer Wins'
                    break
                turn = 'human'

            if isBoardFull(mainBoard):
                winner = 'Tie'
                break

        gameOverSurf = bigFont.render(winner, 1, TEXTCOLOR)
        gameOverRect = gameOverSurf.get_rect()
        gameOverRect.center = (int(WINDOWWIDTH / 2), int(WINDOWHEIGHT / 2))

        newGame = False
        while True:
            if newGame:
                break
            drawBoard(mainBoard)
            windowSurf.blit(gameOverSurf, gameOverRect)
            windowSurf.blit(gameOverSurf2, gameOverRect2)
            pygame.display.update()
            gameClock.tick()
            for event in pygame.event.get():
                if event.type == QUIT:
                    terminate()
                if event.type == MOUSEBUTTONUP:
                    newGame = True
                    break

def makeMove(board, player, column):
    lowest = getLowestFreeSpace(board, column)
    if lowest != -1:
        board[column][lowest] = player


def drawBoard(board, extraToken=None):
    windowSurf.fill(BGCOLOR)

    # draw tokens
    spaceRect = pygame.Rect(0, 0, SPACESIZE, SPACESIZE)
    for x in range(BOARDWIDTH):
        for y in range(BOARDHEIGHT):
            spaceRect.topleft = (XMARGIN + (x * SPACESIZE), YMARGIN + (y * SPACESIZE))
            if board[x][y] == 'red':
                windowSurf.blit(redTokenImg, spaceRect)
            elif board[x][y] == 'black':
                windowSurf.blit(blackTokenImg, spaceRect)

    # draw the extra token
    if extraToken != None:
        if extraToken['color'] == 'red':
            windowSurf.blit(redTokenImg, (extraToken['x'], extraToken['y'], SPACESIZE, SPACESIZE))
        elif extraToken['color'] == 'black':
            windowSurf.blit(blackTokenImg, (extraToken['x'], extraToken['y'], SPACESIZE, SPACESIZE))

    # draw board over the tokens
    for x in range(BOARDWIDTH):
        for y in range(BOARDHEIGHT):
            spaceRect.topleft = (XMARGIN + (x * SPACESIZE), YMARGIN + (y * SPACESIZE))
            windowSurf.blit(boardImg, spaceRect)

    # draw the red and black tokens off to the side
    windowSurf.blit(redTokenImg, redPileRect) # red on the left
    windowSurf.blit(blackTokenImg, blackPileRect) # black on the right



def getNewBoard():
    board = []
    for x in range(BOARDWIDTH):
        board.append([None] * BOARDHEIGHT)
    return board


def getHumanMove(board):
    draggingToken = False
    tokenx, tokeny = None, None
    while True:
        for event in pygame.event.get():
            if event.type == QUIT:
                terminate()
            if not draggingToken and event.type == MOUSEBUTTONDOWN and redPileRect.collidepoint(event.pos):
                # start of dragging on red token pile.
                draggingToken = True
            if draggingToken and event.type == MOUSEMOTION:
                # draw red token being dragged
                tokenx, tokeny = event.pos
            if draggingToken and event.type == MOUSEBUTTONUP:
                # let go of dragging token.

                if tokeny < YMARGIN and tokenx > XMARGIN and tokenx < WINDOWWIDTH - XMARGIN:
                    # let go at the top of the screen.
                    column = int((tokenx - XMARGIN) / SPACESIZE)
                    if isValidMove(board, column):
                        animateDroppingToken(board, column, 'red')
                        board[column][getLowestFreeSpace(board, column)] = 'red'
                        drawBoard(board)
                        pygame.display.update()
                        return
                tokenx, tokeny = None, None
                draggingToken = False
        if tokenx != None and tokeny != None:
            drawBoard(board, {'x':tokenx - int(SPACESIZE / 2), 'y':tokeny - int(SPACESIZE / 2), 'color':'red'})
        else:
            drawBoard(board)

        pygame.display.update()
        gameClock.tick()



def animateDroppingToken(board, column, color):
    x = XMARGIN + column * SPACESIZE
    y = YMARGIN - SPACESIZE
    speed = 1.0

    lowestFreeSpace = getLowestFreeSpace(board, column)

    while True:
        y += int(speed)
        speed += 0.5
        if int((y - YMARGIN) / SPACESIZE) >= lowestFreeSpace:
            return
        drawBoard(board, {'x':x, 'y':y, 'color':color})
        pygame.display.update()
        gameClock.tick()

def animateComputerMoving(board, column):
    x = blackPileRect.left
    y = blackPileRect.top

    speed = 1.0
    # moving the black tile up
    while y > (YMARGIN - SPACESIZE):
        y -= int(speed)
        speed += 0.5
        drawBoard(board, {'x':x, 'y':y, 'color':'black'})
        pygame.display.update()
        gameClock.tick()
    # moving the black tile over
    y = YMARGIN - SPACESIZE
    speed = 1.0
    while x > (XMARGIN + column * SPACESIZE):
        x -= int(speed)
        speed += 0.5
        drawBoard(board, {'x':x, 'y':y, 'color':'black'})
        pygame.display.update()
        gameClock.tick()
    animateDroppingToken(board, column, 'black')


def getComputerMove(board):
    potentialMoves = getPotentialMoves(board, 'black', DIFFICULTY)
    bestMoveScore = max([potentialMoves[i] for i in range(BOARDWIDTH) if isValidMove(board, i)])
    bestMoves = []
    for i in range(len(potentialMoves)):
        if potentialMoves[i] == bestMoveScore:
            bestMoves.append(i)
    return random.choice(bestMoves)


def getPotentialMoves(board, playerTile, lookAhead):
    if lookAhead == 0:
        return [0] * BOARDWIDTH

    potentialMoves = []

    if playerTile == 'red':
        enemyTile = 'black'
    else:
        enemyTile = 'red'

    # Returns (best move, average condition of this state)
    if isBoardFull(board):
        return [0] * BOARDWIDTH

    # Figure out the best move to make.
    potentialMoves = [0] * BOARDWIDTH
    for playerMove in range(BOARDWIDTH):
        dupeBoard = copy.deepcopy(board)
        if not isValidMove(dupeBoard, playerMove):
            continue
        makeMove(dupeBoard, playerTile, playerMove)
        if isWinner(dupeBoard, playerTile):
            potentialMoves[playerMove] = 1
            break
        else:
            # do other player's moves and determine best one
            if isBoardFull(dupeBoard):
                potentialMoves[playerMove] = 0
            else:
                for enemyMove in range(BOARDWIDTH):
                    dupeBoard2 = copy.deepcopy(dupeBoard)
                    if not isValidMove(dupeBoard2, enemyMove):
                        continue
                    makeMove(dupeBoard2, enemyTile, enemyMove)
                    if isWinner(dupeBoard2, enemyTile):
                        potentialMoves[playerMove] = -1
                        break
                    else:
                        results = getPotentialMoves(dupeBoard2, playerTile, lookAhead - 1)
                        potentialMoves[playerMove] += (sum(results) / BOARDWIDTH) / BOARDWIDTH
    return potentialMoves


def getLowestFreeSpace(board, column):
    for y in range(BOARDHEIGHT-1, -1, -1):
        if board[column][y] == None:
            return y
    return -1


def isValidMove(board, move):
    if move < 0 or move >= (BOARDWIDTH):
        return False

    if board[move][0] != None:
        return False

    return True


def isBoardFull(board):
    for x in range(BOARDWIDTH):
        for y in range(BOARDHEIGHT):
            if board[x][y] == None:
                return False
    return True


def isWinner(board, tile):
    # check horizontal spaces
    for y in range(BOARDHEIGHT):
        for x in range(BOARDWIDTH - 3):
            if board[x][y] == tile and board[x+1][y] == tile and board[x+2][y] == tile and board[x+3][y] == tile:
                return True

    # check vertical spaces
    for x in range(BOARDWIDTH):
        for y in range(BOARDHEIGHT - 3):
            if board[x][y] == tile and board[x][y+1] == tile and board[x][y+2] == tile and board[x][y+3] == tile:
                return True

    # check / diagonal spaces
    for x in range(BOARDWIDTH - 3):
        for y in range(3, BOARDHEIGHT):
            if board[x][y] == tile and board[x+1][y-1] == tile and board[x+2][y-2] == tile and board[x+3][y-3] == tile:
                return True

    # check \ diagonal spaces
    for x in range(BOARDWIDTH - 3):
        for y in range(BOARDHEIGHT - 3):
            if board[x][y] == tile and board[x+1][y+1] == tile and board[x+2][y+2] == tile and board[x+3][y+3] == tile:
                return True

    return False

def terminate():
    pygame.quit()
    sys.exit()



if __name__ == '__main__':
    main()