import math
import random
import gym
from gym import spaces, logger
from gym.utils import seeding
import numpy as np
from gym.envs.classic_control import rendering

class Connect4Env(gym.Env):
    metadata = {'render.modes': ['human']}

    def __init__(self):
        self.state = np.zeros([6,7])
        self.viewer = None
        self.difficulty = 3
        self.action_space = spaces.Discrete(7)
        self.reward_1 = 0
        self.reward_2 = 0

        self.seed()

    def setchecker(self, action, color):
        for x in range(len(self.state)):
            if self.state[x, action] == 0:
                self.state[x, action] = color
                return 0.0
        return -10.0

    def checkWin(self, player):
        board = self.state
        boardHeight = len(board[0])
        boardWidth = len(board)
        tile = player
        # 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 step(self, action, color):
    #     assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
    #     done = False
    #     self.reward += self.setchecker(action, color)
    #     if self.checkWin(color):
    #         self.reward = 10.0
    #         done = True
    #         return np.array(self.state), self.reward, done, {}

    #         # print("human", self.checkWin(color))
    #     m = Minimax(self.state)
    #     best_move, value = m.bestMove(self.difficulty, self.state, 2)
    #     self.setchecker(best_move, 2)
    #     # self.render()
    #     # sec = input('Let us wait for user input. Let me know how many seconds to sleep now.\n')
    #     # self.setchecker(int(sec), 2)

    #     if self.checkWin(2):
    #         self.reward += -1.0
    #         done = True
    #         return np.array(self.state), self.reward, done, {}

    #         # print("bot", self.checkWin(2))
    #     return np.array(self.state), self.reward, done, {}

    def step(self, action, color, reward):
        assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
        done = False
        reward += self.setchecker(action, color)
        opponent = 1 if color == 2 else 2 
        if self.checkWin(color):
            reward += 10.0
            done = True
        if self.checkWin(opponent):
            reward -= 1
            done = True
        return np.array(self.state), reward, done, {}




    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]


    def reset(self):
        self.state = np.zeros([6,7])
        self.reward_1 = 0
        self.reward_2 = 0
        return self.state

    def render(self, mode='human'):
        screen_width = 600
        screen_height = 400

        if self.viewer is None:
            self.viewer = rendering.Viewer(screen_width, screen_height)
        for x in range(len(self.state)):
            for y in range(len(self.state[x])):
                circle = rendering.make_circle(20)
                circle.add_attr( rendering.Transform(translation=(y*50 + 140, x*50 + 80)))
                if self.state[x][y] == 1:
                    circle.set_color(0, 1, 0)
                elif self.state[x][y] == 2:
                    circle.set_color(1, 0, 0)
                else:
                    circle.set_color(0, 0, 0)
                self.viewer.add_geom(circle)
        if self.state is None:
            return None

        return self.viewer.render(return_rgb_array=mode == 'rgb_array')

    def close(self):
        if self.viewer:
            self.viewer.close()
            self.viewer = None


class Minimax(object):
    """ Minimax object that takes a current connect four board state
    """
    
    board = None
    colors = [1, 2]
    
    def __init__(self, board):
        # copy the board to self.board
        self.board = np.copy(board)
            
    def bestMove(self, depth, state, curr_player):
        """ Returns the best move (as a column number) and the associated alpha
            Calls search()
        """
        # if np.random.rand() <=0.1:
        #     choices = []
        #     for col in range(7):
        #         if self.isLegalMove(col, state):
        #             choices.append(col)
        #     return random.choices(choices), 1

        # determine opponent's color
        if curr_player == self.colors[0]:
            opp_player = self.colors[1]
        else:
            opp_player = self.colors[0]
        # enumerate all legal moves
        legal_moves = {} # will map legal move states to their alpha values
        for col in range(7):
            # if column i is a legal move...
            if self.isLegalMove(col, state):
                # make the move in column 'col' for curr_player
                temp = self.makeMove(state, col, curr_player)
                legal_moves[col] = -self.search(depth-1, temp, opp_player)
        
        best_alpha = -99999999
        best_move = None
        moves = legal_moves.items()
        random.shuffle(list(moves))
        for move, alpha in moves:
            if alpha >= best_alpha:
                best_alpha = alpha
                best_move = move
        
        return best_move, best_alpha
        
    def search(self, depth, state, curr_player):
        """ Searches the tree at depth 'depth'
            By default, the state is the board, and curr_player is whomever 
            called this search
            
            Returns the alpha value
        """
        
        # enumerate all legal moves from this state
        legal_moves = []
        for i in range(7):
            # if column i is a legal move...
            if self.isLegalMove(i, state):
                # make the move in column i for curr_player
                temp = self.makeMove(state, i, curr_player)
                legal_moves.append(temp)
        
        # if this node (state) is a terminal node or depth == 0...
        if depth == 0 or len(legal_moves) == 0 or self.gameIsOver(state):
            # return the heuristic value of node
            return self.value(state, curr_player)
        
        # determine opponent's color
        if curr_player == self.colors[0]:
            opp_player = self.colors[1]
        else:
            opp_player = self.colors[0]

        alpha = -99999999
        for child in legal_moves:
            # if child == None:
            #     print("child == None (search)")
            alpha = max(alpha, -self.search(depth-1, child, opp_player))
        return alpha

    def isLegalMove(self, column, state):
        """ Boolean function to check if a move (column) is a legal move
        """
        for i in range(6):
            if state[i][column] == 0:
                # once we find the first empty, we know it's a legal move
                return True
        
        # if we get here, the column is full
        return False
    
    def gameIsOver(self, state):
        if self.checkForStreak(state, self.colors[0], 4) >= 1:
            return True
        elif self.checkForStreak(state, self.colors[1], 4) >= 1:
            return True
        else:
            return False
        
    
    def makeMove(self, state, column, color):
        """ Change a state object to reflect a player, denoted by color,
            making a move at column 'column'
            
            Returns a copy of new state array with the added move
        """
        
        temp = np.copy(state)
        for i in range(6):
            if temp[i][column] == 0:
                temp[i][column] = color
                return temp

    def value(self, state, color):
        """ Simple heuristic to evaluate board configurations
            Heuristic is (num of 4-in-a-rows)*99999 + (num of 3-in-a-rows)*100 + 
            (num of 2-in-a-rows)*10 - (num of opponent 4-in-a-rows)*99999 - (num of opponent
            3-in-a-rows)*100 - (num of opponent 2-in-a-rows)*10
        """
        if color == self.colors[0]:
            o_color = self.colors[1]
        else:
            o_color = self.colors[0]
        
        my_fours = self.checkForStreak(state, color, 4)
        my_threes = self.checkForStreak(state, color, 3)
        my_twos = self.checkForStreak(state, color, 2)
        opp_fours = self.checkForStreak(state, o_color, 4)
        #opp_threes = self.checkForStreak(state, o_color, 3)
        #opp_twos = self.checkForStreak(state, o_color, 2)
        if opp_fours > 0:
            return -100000
        else:
            return my_fours*100000 + my_threes*100 + my_twos
            
    def checkForStreak(self, state, color, streak):
        count = 0
        # for each piece in the board...
        for i in range(6):
            for j in range(7):
                # ...that is of the color we're looking for...
                if state[i][j] == color:
                    # check if a vertical streak starts at (i, j)
                    count += self.verticalStreak(i, j, state, streak)
                    
                    # check if a horizontal four-in-a-row starts at (i, j)
                    count += self.horizontalStreak(i, j, state, streak)
                    
                    # check if a diagonal (either way) four-in-a-row starts at (i, j)
                    count += self.diagonalCheck(i, j, state, streak)
        # return the sum of streaks of length 'streak'
        return count
            
    def verticalStreak(self, row, col, state, streak):
        consecutiveCount = 0
        for i in range(row, 6):
            if state[i][col] == state[row][col]:
                consecutiveCount += 1
            else:
                break
    
        if consecutiveCount >= streak:
            return 1
        else:
            return 0
    
    def horizontalStreak(self, row, col, state, streak):
        consecutiveCount = 0
        for j in range(col, 7):
            if state[row][j] == state[row][col]:
                consecutiveCount += 1
            else:
                break

        if consecutiveCount >= streak:
            return 1
        else:
            return 0
    
    def diagonalCheck(self, row, col, state, streak):

        total = 0
        # check for diagonals with positive slope
        consecutiveCount = 0
        j = col
        for i in range(row, 6):
            if j > 6:
                break
            elif state[i][j] == state[row][col]:
                consecutiveCount += 1
            else:
                break
            j += 1 # increment column when row is incremented
            
        if consecutiveCount >= streak:
            total += 1

        # check for diagonals with negative slope
        consecutiveCount = 0
        j = col
        for i in range(row, -1, -1):
            if j > 6:
                break
            elif state[i][j] == state[row][col]:
                consecutiveCount += 1
            else:
                break
            j += 1 # increment column when row is incremented

        if consecutiveCount >= streak:
            total += 1

        return total
