connect 4

connect4.py

import numpy as np
import sys
import math


class ConnectFour:
	def __init__(self, row_count, column_count):
		self.row_count = row_count
		self.column_count = column_count

	#private
	def create_board(self):
		board = np.zeros((self.row_count, self.column_count))
		return board

	#private
	def drop_piece(self, board, row, column, piece):
		board[row][column] = piece

	#private
	def is_valid_location(self, board, column):
		return board[self.row_count-1][column] == 0

	#private
	def get_next_open_row(self, board, column):
		for r in range(self.row_count):
			if board[r][column] == 0:
				return r
	
	#private
	def winning_move(self, board, piece):
		for c in range(self.column_count-3):
			for r in range(self.row_count):
				if board[r][c] == piece and board[r][c+1] == piece and board[r][c+2] == piece and board[r][c+3] == piece:
					return True
		for c in range(self.column_count):
			for r in range(self.row_count-3):
				if board[r][c] == piece and board[r+1][c] == piece and board[r+2][c] == piece and board[r+3][c] == piece:
					return True
		for c in range(self.column_count-3):
			for r in range(self.row_count-3):
				if board[r][c] == piece and board[r+1][c+1] == piece and board[r+2][c+2] == piece and board[r+3][c+3] == piece:
					return True
		for c in range(self.column_count-3):
			for r in range(3, self.row_count):
				if board[r][c] == piece and board[r-1][c+1] == piece and board[r-2][c+2] == piece and board[r-3][c+3] == piece:
					return True
	#private
	def play(self, piece, board):
		column = int(input("Player {}'s turn (1-7)".format(piece)))
		column -= 1
		if board[self.row_count-1][column] == 0:
			row = self.get_next_open_row(board, column)
			self.drop_piece(board, row, column, piece)
			if self.winning_move(board, piece):
				print("~"*15 + "Player {} wins".format(piece) + "~"*15)
				print(np.flip(board, 0))
				print("~"*15 + "Player {} wins".format(piece) + "~"*15)
				sys.exit()

	#public
	def main(self):
		board = self.create_board()
		print(np.flip(board, 0))
		game_over = False
		turn = 0
		while not game_over:
			if turn == 0:
				self.play(1, board)
			else: 
				self.play(2, board)
			print(np.flip(board, 0))	
			turn += 1
			turn = turn % 2


if __name__ == "__main__":
	game = ConnectFour(6,7)
	game.main()