import sys

import pygame
import numpy as np
from matplotlib import pyplot as plt

from reinforcementLearning.ReinforcementLearning import ReinforcementLearning
from utilities.constants import WIDTH, HEIGHT, SQUARE_SIZE, WHITE, GREEN
from utilities.gameManager import GameManager
from minimax.minimaxAlgo import MiniMax

FPS = 60
WIN = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Draughts")


def getRowColFromMouse(pos: dict) -> tuple:
    """
    Gets the row and column from the mouse position
    :param pos: X and Y position of the mouse
    :return: Row and column
    """
    x, y = pos
    row = y // SQUARE_SIZE
    col = x // SQUARE_SIZE
    return row, col


def drawText(text: str, font:  pygame.font.SysFont, colour: tuple, surface: pygame.display, x: float, y: int) -> None:
    """
    Draws text on the screen
    :param text: Text to draw
    :param font: System font
    :param colour: Colour of the text
    :param surface: The display surface
    :param x: X position of the text
    :param y: Y position of the text
    :return None
    """
    textobj = font.render(text, 1, colour)
    textrect = textobj.get_rect()
    textrect.topleft = (x, y)
    surface.blit(textobj, textrect)


def drawMultiLineText(surface: pygame.display, text: str, pos: dict, font: pygame.font.SysFont, colour: tuple = pygame.Color('black')) -> None:
    """
    Draws multiline text on the screen
    :param surface: the display surface
    :param text: text to draw
    :param pos: X and Y position of the text
    :param font: System font
    :param colour: colour of the text
    :return None
    """
    words = [word.split(' ') for word in text.splitlines()]  # 2D array where each row is a list of words.
    space = font.size(' ')[0]  # The width of a space.
    max_width, max_height = surface.get_size()
    x, y = pos
    word_height = None
    for line in words:
        for word in line:
            word_surface = font.render(word, 0, colour)
            word_width, word_height = word_surface.get_size()
            if x + word_width >= max_width:
                x = pos[0]  # Reset the x.
                y += word_height  # Start on new row.
            surface.blit(word_surface, (x, y))
            x += word_width + space
        x = pos[0]  # Reset the x.
        y += word_height  # Start on new row.


def main(difficulty: int = 0) -> None:
    """
    Main function, that shows the menu before running the game
    :param difficulty: difficulty of minimax
    :return: None
    """
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    menuClock = pygame.time.Clock()
    click = False
    width = screen.get_width()
    font = pygame.font.SysFont("", 25)

    if difficulty == 0:
        while True:
            # menu
            screen.fill((128, 128, 128))
            drawText('Main Menu', font, (255, 255, 255), screen, width / 2, 20)

            mx, my = pygame.mouse.get_pos()

            easy = pygame.Rect(width / 2 - 50, 100, 200, 50)
            pygame.draw.rect(screen, (0, 255, 0), easy)
            drawText("easy", font, (255, 255, 255), screen, width / 2, 100)
            medium = pygame.Rect(width / 2 - 50, 200, 200, 50)
            pygame.draw.rect(screen, (255, 125, 0), medium)
            drawText("medium", font, (255, 255, 255), screen, width / 2, 200)
            hard = pygame.Rect(width / 2 - 50, 300, 200, 50)
            pygame.draw.rect(screen, (255, 0, 0), hard)
            drawText("hard", font, (255, 255, 255), screen, width / 2, 300)
            rules = pygame.Rect(width / 2 - 50, 400, 200, 50)
            pygame.draw.rect(screen, (0, 0, 255), rules)
            drawText("rules", font, (255, 255, 255), screen, width / 2, 400)
            quitGame = pygame.Rect(width / 2 - 50, 500, 200, 50)
            pygame.draw.rect(screen, (0, 0, 0), quitGame)
            drawText("quit", font, (255, 255, 255), screen, width / 2, 500)

            if easy.collidepoint((mx, my)):
                if click:
                    difficulty = 1
                    break
            if medium.collidepoint((mx, my)):
                if click:
                    difficulty = 3
                    break
            if hard.collidepoint((mx, my)):
                if click:
                    difficulty = 5
                    break
            if rules.collidepoint((mx, my)):
                if click:
                    rulesGUI()
                    break
            if quitGame.collidepoint((mx, my)):
                if click:
                    pygame.quit()
                    sys.exit()
            click = False
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    pygame.quit()
                    sys.exit()
                if event.type == pygame.MOUSEBUTTONDOWN:
                    if event.button == 1:
                        click = True

            pygame.display.update()
            menuClock.tick(60)

    game(difficulty)


def rulesGUI() -> None:
    """
    Shows the rules of the game
    :return: None
    """
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    menuClock = pygame.time.Clock()
    click = False
    width = screen.get_width()
    titleFont = pygame.font.SysFont("", 48)
    font = pygame.font.SysFont("", 21)
    while True:
        screen.fill((128, 128, 128))
        drawText("Rules", titleFont, (255, 255, 255), screen, width / 2, 20)

        mx, my = pygame.mouse.get_pos()
        drawMultiLineText(screen, """Both the player and AI start with 12 pieces on the dark squares of the three rows closest to that 
        player's side. The row closest to each player is called the kings row or crownhead. The player moves first. 
        Then turns alternate. 
\n
Move rules
\n
There are two different ways to move in  utilities:
\n
Simple move: A simple move consists of moving a piece one square diagonally to an adjacent unoccupied dark square. 
Uncrowned pieces can move diagonally forward only; kings can move in any diagonal direction. Jump: A jump consists of 
moving a piece that is diagonally adjacent an opponent's piece, to an empty square immediately beyond it in the same 
direction (thus "jumping over" the opponent's piece front and back ). Pieces can jump diagonally forward only; kings 
can jump in any diagonal direction. A jumped piece is considered "captured" and removed from the game. Any piece, 
king or piece, can jump a king.
\n
Forced capture, is always mandatory: if a player has the option to jump, he/she must take it, even if doing so 
results in disadvantage for the jumping player. For example, a piecedated single jump might set up the player such 
that the opponent has a multi-jump in reply.
\n
Multiple jumps are possible, if after one jump, another piece is immediately eligible to be jumped by the moved 
piece—even if that jump is in a different diagonal direction. If more than one multi-jump is available, the player 
can choose which piece to jump with, and which sequence of jumps to make. The sequence chosen is not required to be 
the one that maximizes the number of jumps in the turn; however, a player must make all available jumps in the 
sequence chosen. Kings If a piece moves into the kings row on the opponent's side of the board, it is crowned as a 
king and gains the ability to move both forward and backward. If a piece moves into the kings row or if it jumps into 
the kings row, the current move terminates; the piece is crowned as a king but cannot jump back out as in a 
multi-jump until the next move.""", (50, 50), font)
        back = pygame.Rect(width / 2 - 50, 700, 200, 50)
        pygame.draw.rect(screen, (0, 0, 0), back)
        drawText("back", font, (255, 255, 255), screen, width / 2, 700)

        if back.collidepoint((mx, my)):
            if click:
                main()
                break

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:
                    click = True

        pygame.display.update()
        menuClock.tick(60)


def game(difficulty: int) -> None:
    """
    Runs the game with the given difficulty. Used for training and testing the RL algorithm
    :param difficulty: The difficulty of the minimax algorithm
    """
    run = True
    clock = pygame.time.Clock()
    gameManager = GameManager(WIN, GREEN)
    rl = ReinforcementLearning(gameManager.board.getAllMoves(WHITE), gameManager.board, WHITE, gameManager)
    # model = rl.buildMainModel()
    rl.model.load_weights("./modelWeights/model_final.h5")
    mm = MiniMax()
    totalReward = []
    winners = []
    for i in range(50):
        score = 0
        for j in range(200):
            print(j)
            clock.tick(FPS)
            reward = 0
            if gameManager.turn == WHITE:
                # mm = MiniMax()
                # value, newBoard = mm.AI(difficulty, WHITE, gameManager)
                # gameManager.aiMove(newBoard)
                # reward, newBoard = rl.AITrain(gameManager.board)
                newBoard = rl.AITest(gameManager.board)

                if newBoard is None:
                    print("Cannot make move")
                    continue
                gameManager.aiMove(newBoard)

            gameManager.update()
            pygame.display.update()

            if gameManager.turn == GREEN:
                value, newBoard = mm.AI(difficulty, GREEN, gameManager)
                gameManager.aiMove(newBoard)

            score += reward

            if gameManager.winner() is not None:
                print("Green" if gameManager.winner() == GREEN else "White", " wins")
                # with open(f"winners-{difficulty}.txt", "a+") as f:
                #     f.write(str(gameManager.winner()) + "\n")
                winners.append(gameManager.winner())
                break

            # for event in pygame.event.get():
            #     if event.type == pygame.QUIT:
            #         break
            #     if event.type == pygame.MOUSEBUTTONDOWN:
            #         pos = pygame.mouse.get_pos()
            #         row, col = getRowColFromMouse(pos)
            #         # if gameManager.turn == GREEN:
            #         gameManager.select(row, col)

            gameManager.update()
            pygame.display.update()

        if gameManager.winner() is None:
            # with open(f"winners-{difficulty}.txt", "a+") as f:
            #     f.write(str(0) + "\n")
            winners.append(0)
        gameManager.reset()
        rl.resetScore()
        print("Game: ", i, " Reward: ", score)
        # with open(f"rewards-{difficulty}.txt", "a+") as f:
        #     f.write(str(score) + "\n")

        totalReward.append(score)
        # save model weights every 25 games
        # if i % 250 == 0 and i != 0:
        #     rl.model.save("./modelWeights/model_" + str(i) + ".h5")
    # pygame.quit()

    # rl.model.save("./modelWeights/model_final.h5")
    change_in_rewards = [0]  # Initialize with 0 for the first episode
    for i in range(1, len(totalReward)):
        change_in_reward = totalReward[i] - totalReward[i - 1]
        change_in_rewards.append(change_in_reward)

    # with open(f"changeInRewards-{difficulty}.txt", "a+") as f:
    #     for i in change_in_rewards:
    #         f.write(str(i) + "\n")

    # episodes = list(range(1, len(totalReward) + 1))
    #
    # plt.plot(episodes, change_in_rewards)
    # plt.xlabel('Training Games')
    # plt.ylabel('Change in Game Reward')
    # plt.title('Change in Game Reward vs. Training Games')
    # plt.grid(True)
    # plt.show()
    #
    # plt.plot([i for i in range(len(totalReward))], totalReward)
    # plt.xlabel("Games")
    # plt.ylabel("Reward")
    # plt.show()

    fig, ax = plt.subplots()
    bar = ax.bar(["Draw", "White", "Green"], [winners.count(0), winners.count(WHITE), winners.count(GREEN)])
    ax.set(xlabel='Winner', ylabel='Frequency', ylim=[0, 500])
    ax.set_title(f"Winners for difficulty — {difficulty}")
    ax.bar_label(bar)
    plt.show()


# difficulties = [3, 5, 7, 9]
#
# for diff in difficulties:
#     main(diff)
main(3)