import numpy as np 
import pyinputplus as pyip 
import pygame, math, sys

BLUE = (0,0,255)
BLACK = (0,0,0)
RED = (255,0,0)
YELLOW = (255,255,0)

ROW_COUNT = 6
COL_COUNT = 7

def create_board():
	board = np.zeros((ROW_COUNT,COL_COUNT))
	return board

def drop_piece(board, row, col, piece):
	board[row][col] = piece

def is_valid_location(board, col):
	if board[ROW_COUNT-1][col] == 0:
		return True
	if board[ROW_COUNT-1][0] != 0 and board[ROW_COUNT-1][1] != 0 and \
	board[ROW_COUNT-1][2] != 0 and board[ROW_COUNT-1][3] != 0 and \
	board[ROW_COUNT-1][4] != 0 and board[ROW_COUNT-1][5] != 0 and \
	board[ROW_COUNT-1][6] != 0:
		return False

def get_next_open_row(board, col):
	for r in range(ROW_COUNT):
		if board[r][col] == 0:
			return r

def print_board(board):
	print(np.flip(board, 0))

def winning_move(board,piece):
	# Check horizontal location for win
	for c in range(COL_COUNT-3):
		for r in range(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

	# Check vertical location for win
	for c in range(COL_COUNT):
		for r in range(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

	# Check for positively slope diagonals
	for c in range(COL_COUNT-3):
		for r in range(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

	# Check for negatively slope diagonals
	for c in range(COL_COUNT-3):
		for r in range(3,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

def draw_board(board):
	for c in range(COL_COUNT):
		for r in range(ROW_COUNT):
			pygame.draw.rect(screen, BLUE, 
				(c*SQUARESIZE, r*SQUARESIZE+SQUARESIZE, SQUARESIZE, SQUARESIZE))
			pygame.draw.circle(screen, BLACK, (int(c*SQUARESIZE+SQUARESIZE/2), 
				int(r*SQUARESIZE+SQUARESIZE+SQUARESIZE/2)), RADIUS)

	for c in range(COL_COUNT):
		for r in range(ROW_COUNT):
			if board[r][c] == 1:
				pygame.draw.circle(screen, RED, (int(c*SQUARESIZE+SQUARESIZE/2), 
					height - int(r*SQUARESIZE+SQUARESIZE/2)), RADIUS)
			elif board[r][c] == 2:
				pygame.draw.circle(screen, YELLOW, (int(c*SQUARESIZE+SQUARESIZE/2), 
					height - int(r*SQUARESIZE+SQUARESIZE/2)), RADIUS)

	pygame.display.update()

board = create_board()
print_board(board)
game_over = False
turn = 0

pygame.init()

SQUARESIZE = 100

width = COL_COUNT * SQUARESIZE
height = (ROW_COUNT + 1) * SQUARESIZE
size = (width, height)

RADIUS = int(SQUARESIZE/2 - 5)

screen = pygame.display.set_mode(size)
draw_board(board)
pygame.display.update()

myfont = pygame.font.SysFont('monospace', 75)

while not game_over:

	for event in pygame.event.get():
		if event.type == pygame.QUIT:
			sys.exit()

		if event.type == pygame.MOUSEMOTION:
			pygame.draw.rect(screen, BLACK, (0,0,width, SQUARESIZE))
			posx = event.pos[0]
			if turn == 0:
				pygame.draw.circle(screen, RED, (posx, int(SQUARESIZE/2)), RADIUS)

			else:
				pygame.draw.circle(screen, YELLOW, (posx, int(SQUARESIZE/2)), RADIUS)

			pygame.display.update()

		if event.type == pygame.MOUSEBUTTONDOWN:
			pygame.draw.rect(screen, BLACK, (0,0,width, SQUARESIZE))
			# Ask for Player 1 input
			if turn == 0:
				posx = event.pos[0]
				col = int(math.floor(posx/SQUARESIZE))
				
				if is_valid_location(board, col):
					row = get_next_open_row(board, col)
					drop_piece(board, row, col, 1)

					if winning_move(board,1):
						label = myfont.render('PLAYER 1 WINS!', 1, RED)
						screen.blit(label, (40,10))
						game_over = True

				if is_valid_location(board, col) == False:
					label = myfont.render('TIE GAME!', 1, BLUE)
					screen.blit(label, (150,10))
					game_over = True


			# Ask for Player 2 input
			else:
				posx = event.pos[0]
				col = int(math.floor(posx/SQUARESIZE))
				
				if is_valid_location(board, col):
					row = get_next_open_row(board, col)
					drop_piece(board, row, col, 2)

					if winning_move(board,2):
						label = myfont.render('PLAYER 2 WINS!', 1, YELLOW)
						screen.blit(label, (40,10))
						game_over = True

				if is_valid_location(board, col) == False:
					label = myfont.render('TIE GAME!', 1, BLUE)
					screen.blit(label, (150,10))
					game_over = True

			print_board(board)
			draw_board(board)

			turn += 1
			turn = turn % 2

			if game_over:
				pygame.time.wait(1000)


