import sys
import pygame
import numpy as np
import random
import copy
import time

from constants import *

# Initialize pygame
pygame.init()
screen = pygame.display.set_mode((WIDTH, LENGHT))
pygame.display.set_caption("Tic Tac Toe Game - AI ")
screen.fill( BACKGROUND_COLOR )

#Class : Matrix
class Matrix:
    def __init__(self):
        self.cases = np.zeros((GRID_ROWS, GRID_COLS))
        self.empty_cases = self.cases
        self.marked_case = 0


    def winer(self, show=False):
        '''
        Check if there is a winner
        return 0 if there is no winner yet
        return 1 if player 1 win
        return 2 if player 2 win
        '''


        # Check vertical lines
        for column in range(GRID_COLS):
            if self.cases[0][column] == self.cases[1][column] == self.cases[2][column] != 0:
                if show:
                    color = PLAYER1_COLOR if self.cases[0][column] == 1 else PLAYER2_COLOR
                    initial_pos = (column * CASE_SIZE + CASE_SIZE//2, 20)
                    final_pos = (column * CASE_SIZE + CASE_SIZE//2, LENGHT - 20)
                    pygame.draw.line(screen, color, initial_pos, final_pos, LINE_WIDTH)
                return self.cases[0][column]
            
        # Check horizontal lines
        for row in range(GRID_ROWS):
            if self.cases[row][0] == self.cases[row][1] == self.cases[row][2] != 0:
                if show:
                    color = PLAYER1_COLOR if self.cases[0][column] == 1 else PLAYER2_COLOR
                    initial_pos = (20, row * CASE_SIZE + CASE_SIZE//2)
                    final_pos = (WIDTH - 20, row * CASE_SIZE + CASE_SIZE//2)
                    pygame.draw.line(screen, color, initial_pos, final_pos, LINE_WIDTH)
                return self.cases[row][0]
       
        # Check diagonals
        if self.cases[0][0] == self.cases[1][1] == self.cases[2][2] != 0:
            if show:
                    color = PLAYER1_COLOR if self.cases[0][column] == 1 else PLAYER2_COLOR
                    initial_pos = (20, 20)
                    final_pos = (WIDTH - 20, LENGHT - 20)
                    pygame.draw.line(screen, color, initial_pos, final_pos, LINE_WIDTH)
            return self.cases[0][0]
        if self.cases[0][2] == self.cases[1][1] == self.cases[2][0] != 0:
            if show:
                    color = PLAYER1_COLOR if self.cases[0][column] == 1 else PLAYER2_COLOR
                    initial_pos = (20, LENGHT - 20)
                    final_pos = (WIDTH - 20, 20)
                    pygame.draw.line(screen, color, initial_pos, final_pos, LINE_WIDTH)
            return self.cases[0][2]
        # Check if there is no winner yet
        return 0
    
    def empty_case(self, row,column):
        return self.cases[row][column] == 0
    
    def mark_case(self, row, column, player):
        self.cases[row][column] = player
        self.marked_case += 1

    def get_empty_cases(self):
        empty_cases = []
        for row in range(GRID_ROWS):
            for column in range(GRID_COLS):
                if self.empty_case(row, column):
                    empty_cases.append((row, column))
        return empty_cases
    
    def isfull(self):
        return self.marked_case == 9
    
    def isempty(self):
        return self.marked_case == 0

#Class : AI
class AI:

    def __init__(self, level= 1, player=2):
        self.level = level
        self.player = player

    def randomM(self,matrix):
        empty_cases = matrix.get_empty_cases()
        a = random.randrange(0,len(empty_cases))

        return empty_cases[a] # (row, column)
    
    def minimax(self, matrix, maximizing):
        
        #Check if terminal case
        case = matrix.winer()
        #player 1 win
        if case == 1:
            return 1, None # eval, move
        #player 2 win
        if case == 2:
            return -1, None # eval, move
        #draw
        elif matrix.isfull():
            return 0, None # eval, move
        
        if maximizing:
            max_eval = -100
            best_move = None
            empty_cases = matrix.get_empty_cases()

            for (row,column) in empty_cases:
                temp_matrix = copy.deepcopy(matrix)
                temp_matrix.mark_case(row, column, 1)
                eval = self.minimax(temp_matrix, False)[0]
                if eval > max_eval:
                    max_eval = eval
                    best_move = (row, column)

            return max_eval, best_move # eval, move
        
        elif not maximizing:
            min_eval = 100
            best_move = None
            empty_cases = matrix.get_empty_cases()

            for (row,column) in empty_cases:
                temp_matrix = copy.deepcopy(matrix)
                temp_matrix.mark_case(row, column, self.player)
                eval = self.minimax(temp_matrix, True)[0]
                if eval < min_eval:
                    min_eval = eval
                    best_move = (row, column)

            return min_eval, best_move # eval, move


    def evaluate(self, main_matrix):
        if self.level == 0:
            #random choice
            eval = 'random'
            move = self.randomM(main_matrix)
        else:
            #minimax algo choice
            print("AI is thinking...")
            time.sleep(1)
            eval, move = self.minimax(main_matrix, False)

        print(f"AI has chosen the pos {move} with an eval of {eval}")

        return move # (row, column)

#Class : Game 
class Game:
    def __init__(self):
        self.matrix = Matrix()
        self.ai = AI()
        self.player = 1 #1 = X, 2 = O
        self.gamemode = 'ai' #1v1 or 1vAI
        self.running = True
        self.ai_starts = False
        self.lines()

    def draw_figures(self, row, column):
        if self.player == 1:
            #draw X
            start_desc = (column * CASE_SIZE + CASE_SIZE // 4, row * CASE_SIZE + CASE_SIZE // 4)
            end_desc = (column * CASE_SIZE + CASE_SIZE * 3 // 4, row * CASE_SIZE + CASE_SIZE * 3 // 4) 
            pygame.draw.line(screen,PLAYER1_COLOR, start_desc, end_desc, CROSS_WIDTH)
            start_asc = (column * CASE_SIZE + CASE_SIZE // 4, row * CASE_SIZE + CASE_SIZE * 3 // 4)
            end_asc = (column * CASE_SIZE + CASE_SIZE * 3 // 4, row * CASE_SIZE + CASE_SIZE // 4)
            pygame.draw.line(screen,PLAYER1_COLOR, start_asc, end_asc, CROSS_WIDTH)
        elif self.player == 2:
            #draw O
            center = ( column * CASE_SIZE + CASE_SIZE // 2, row * CASE_SIZE + CASE_SIZE // 2)
            pygame.draw.circle(screen, PLAYER2_COLOR,center, CIRCLE_RADIUS, CIRCLE_WIDTH)

    def make_move(self, row, col):
        self.matrix.mark_case(row, col, self.player)
        self.draw_figures(row, col)
        self.player_turn()

    def player_turn(self):
        self.player = self.player % 2 + 1

    def lines(self):
        # Horizontal lines
        pygame.draw.line(screen, LINE_COLOR, (0, CASE_SIZE), (WIDTH,CASE_SIZE), LINE_WIDTH)
        pygame.draw.line(screen, LINE_COLOR, (0, LENGHT - CASE_SIZE), (WIDTH, LENGHT - CASE_SIZE), LINE_WIDTH)

        #Vertical lines
        pygame.draw.line(screen, LINE_COLOR, (CASE_SIZE, 0), (CASE_SIZE, LENGHT), LINE_WIDTH)
        pygame.draw.line(screen, LINE_COLOR, (WIDTH - CASE_SIZE, 0), (WIDTH - CASE_SIZE, LENGHT), LINE_WIDTH)

    def change_gamemode(self):
        if self.gamemode == 'ai':
            self.gamemode = '1v1'
        else:
            self.gamemode = 'ai'

    def isover(self):
        return self.matrix.winer(show=True) != 0 or self.matrix.isfull()
    
    def reset(self):
        self.__init__()
        screen.fill(BACKGROUND_COLOR)
        self.lines()
        if self.ai_starts:
            self.ai_first_move()

    def ai_first_move(self):
        if self.gamemode == 'ai' and self.player == self.ai.player and self.ai_starts:
            row, column = self.ai.evaluate(self.matrix)
            self.make_move(row, column)
            self.ai_starts = False

    # ...
    def start_ai_turn(self):
        if self.gamemode == 'ai' and self.player == self.ai.player and self.ai_starts:
            row, column = self.ai.evaluate(self.matrix)
            self.make_move(row, column)
            self.ai_starts = False
    # ...



def main():

    # Object of Game class
    game = Game()
    matrix = game.matrix
    ai = game.ai

    # Main loop
    while True:

        for event in pygame.event.get():

            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

            if event.type == pygame.KEYDOWN:
                
                # g-gamemode
                if event.key == pygame.K_g:
                    game.change_gamemode()
                
                # s-start
                if event.key == pygame.K_s:
                    game.ai_starts = not game.ai_starts
                    game.start_ai_turn()

                #r-restart
                if event.key == pygame.K_r:
                    game.reset()
                    matrix = game.matrix
                    ai = game.ai
                    
                # 0-random ai
                if event.key == pygame.K_0:
                    ai.level = 0

                # 1-easy ai
                if event.key == pygame.K_1:
                    ai.level = 1
                
            #click event
            if event.type == pygame.MOUSEBUTTONDOWN:
                pos = event.pos
                row = pos[1] // CASE_SIZE
                column = pos[0] // CASE_SIZE
                
                if matrix.empty_case(row,column) and game.running:
                    game.make_move(row, column)

                    if game.isover():
                        game.running = False

        if game.gamemode == 'ai' and game.player == ai.player and game.running:
            # update the screen
            pygame.display.update()

            # AI turn
            row, column = ai.evaluate(matrix)
            game.make_move(row, column)

            if game.isover():
                game.running = False

        pygame.display.update()
main()