import time
import sys
import curses 
import random
from copy import deepcopy

inf = 2**31

class C4(object):
    def __init__(self, width=7, height=6):
        self.width = width
        self.height = height
        self.grid = [[' ' for y in range(self.height)] for x in range(self.width)]
        self.turn = 'O'
        self.cpu_move('O')

    def legal_moves(self):
        moves = []
        for x in range(self.width):
            if self.grid[x][0] == ' ':
                moves.append(x)
        return moves

    def output(self, w):
        for x in range(self.width):
            if x in self.legal_moves():
                w.addstr(2, 6+x*4, str(x+1))
            for y in range(self.height):
                px, py = 6+x*4, 4+y*2
                for dx,dy in [(-2,-1),(2,-1),(-2,1),(2,1)]:
                    w.addch(py+dy, px+dx, curses.ACS_SSSS)

                for dx in (-1, 0, 1):
                    for dy in (-1, 1):
                        w.addch(py+dy, px+dx, curses.ACS_HLINE)
                for dx in (-2, 2):
                    w.addch(py, px+dx, curses.ACS_VLINE)
                w.addstr(py, px, self.grid[x][y])

        win = self.win()
        if win is not None:
            player, discs = win
            for x, y in discs:
                px, py = 6+x*4, 4+y*2
                w.addstr(py, px, self.grid[x][y], curses.A_BOLD)

        w.addstr(4, 40, 'You\'re X')
        w.addstr(5, 40, 'CPU is O')
        if win is None:
            w.addstr(7, 40, '1-7) place a disc')
            w.addstr(8, 40, 'a)   let CPU place a disc for you')
            w.addstr(9, 40, 'r)   restart')
            w.addstr(11, 40, 'It\'s {} turn'.format({'X': 'your', 'O': 'CPU\'s'}[self.turn]))
        else:
            w.addstr(7, 40, '{} won'.format({'X': 'You', 'O': 'CPU'}[win[0]]))
            w.addstr(9, 40, 'r) restart')

    def move(self, move, player='X'):
        if self.turn != player:
            return
        if move not in self.legal_moves():
            return
        if self.terminal():
            return
        for y in range(self.height-1, -1, -1):
            if self.grid[move][y] == ' ':
                self.grid[move][y] = player
                self.turn = {'X': 'O', 'O': 'X'}[player]
                return

    def cpu_move(self, player):
        if self.terminal():
            return
        moves = alphabeta(deepcopy(self), 5, -inf, +inf, player)
        self.move(random.choice(moves), player)

    def successors(self):
        ss = []
        for move in self.legal_moves():
          s = deepcopy(self)
          s.move(move, self.turn)
          ss.append((move, s))
        return ss

    def terminal(self):
        return len(self.legal_moves()) == 0 or self.win() is not None

    def utility(self, player):
        win = self.win()
        if win is None:
            return 0
        if win[0] == player:
            return 1
        return -1

    def win(self):
        for x in range(self.width):
            for y in range(self.height):
                for dx,dy in [(0,1), (1,0), (1,1)]:
                    cx, cy = x, y
                    player = self.grid[cx][cy]
                    if player == ' ':
                        continue
                    win = True
                    for i in range(3):
                        cx += dx
                        cy += dy
                        try:
                            if self.grid[cx][cy] != player:
                                win = False
                                break
                        except IndexError:
                            win = False
                            break
                    if win:
                        return player, [(x+dx*i,y+dy*i) for i in range(4)]
        return None


def alphabeta(node, depth, alpha, beta, player, fn=None):
    out = False
    if fn is None:
        out = True
        fn = 'max'
    if depth == 0 or node.terminal():
        a = node.utility(player)
        return a
    if fn == 'max':
        moves = []
        for succ_k, succ in node.successors():
            alpha_ = alphabeta(succ, depth-1, alpha, beta, player, 'min')
            if alpha_ > alpha:
                alpha = alpha_
                moves = [succ_k]
            elif alpha_ == alpha:
                moves.append(succ_k)
            if beta <= alpha:
                break
        if out:
            return moves
        else:
            return alpha
    else:
        for succ_k, succ in node.successors():
            beta = min(beta, alphabeta(succ, depth-1, alpha, beta, player, 'max'))
            if beta <= alpha:
                break
        return beta
            
    

def main():
    screen = curses.initscr()
    curses.noecho() 
    curses.curs_set(0) 
    screen.keypad(1) 

    c4 = C4()
    while True:
        screen.clear()
        screen.border()
        c4.output(screen)
        event = screen.getch()
        if ord('1') <= event <= ord('7'):
            c4.move(event-ord('1'))
            c4.cpu_move('O')
        if event == ord('a'):
            c4.cpu_move('X')
            c4.cpu_move('O')
        if event == ord('r'):
            c4 = C4()

try:
    main()
finally:
    curses.endwin()