#include "game.hpp"

#include <chrono>
#include <vector>
#include <algorithm>
#include <cstdlib>
#include <limits>
#include <thread>

//Used to calculate ticks of the game
uint64_t time_milli() {
    using namespace std::chrono;
    return duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
}

//Default constructor without a FEN
Game::Game(SDL_Window* window, SDL_Surface* surface)
{
    this->window = window;
    this->surface = surface;
    running = false;
    holding_k = false;
}

//Construct with specific board
Game::Game(SDL_Window* window, SDL_Surface* surface, std::string board_fen)
{
    this->window = window;
    this->surface = surface;
    board = Board(board_fen);
    running = false;
    holding_k = false;
}

//Main loop
void Game::run()
{
    //calculates elapsed quantum and runs a tick, always draws
    uint64_t last_update_time = time_milli(), delta_time;
    uint64_t time_quantum = 1000/60;
    running = true;
    while(running)
    {
        delta_time = time_milli()-last_update_time;
        while(delta_time > time_quantum && running)
        {
            tick();

            delta_time -= time_quantum;
            last_update_time += time_quantum;
        }
        draw();
    }
}

void Game::tick()
{
    static bool ai_lost = false;
    //Check for user interface
    SDL_Event e;
    while(running && SDL_PollEvent(&e) != 0)
    {   
        //Pressing q or hitting close should stop the game
        if (e.type == SDL_QUIT || (e.type == SDL_KEYUP && e.key.keysym.sym == SDLK_q))
        {
            running = false;
        }
        //holding K needs to be kept track of for pawn promotion between queen or knight
        else if (e.type == SDL_KEYUP && e.key.keysym.sym == SDLK_k)
        {
            holding_k = false;
        }
        else if (e.type == SDL_KEYDOWN && e.key.keysym.sym == SDLK_k)
        {
            holding_k = true;
        }
        //Clicking causes the selection of spaces
        else if (e.type == SDL_MOUSEBUTTONDOWN)
        {
            int x,y;
            SDL_GetMouseState(&x, &y);
            x/=SPRITE_SIZE;
            y/=SPRITE_SIZE;
            process_click(x,y);
            //Draw now because the AI turn takes a fair amount of time
            draw();
        }
    }

    if (!ai_lost && BLACK == board.get_active_color())
    {
        if (-1 == do_ai_move())
        {
            ai_lost = true;
        }
    }
}

void Game::draw()
{
    board.draw_board(surface);

    SDL_UpdateWindowSurface(window);
}

//Processes the click on the board
void Game::process_click(int x, int y)
{
    //Only allow clicks if its white's turn
    if (WHITE == board.get_active_color())
    {
        //Get the current space and target space
        int current_selected_space = board.get_selected_space();
        int target_selected_space = x+y*BOARD_SIZE;
        if (current_selected_space == -1)
        {
            //Nothing selected before, select this tile
            board.set_selected_space(x,y);
        }
        else
        {
            //Get the target spaces for the previously selected space
            std::vector<int> target_spaces = board.get_moves_for_space(current_selected_space, true);

            //Move if the mouse clicked a possible target, otherwise change the selected space to the new one
            if (target_spaces.end() != std::find(target_spaces.begin(), target_spaces.end(), target_selected_space))
            {
                board.do_move(current_selected_space, target_selected_space, holding_k);
            }
            else
            {
                board.set_selected_space(target_selected_space);
            }
        }
    }
}

int Game::do_ai_move()
{
    //Thread management
    std::vector<std::thread> threads;
    std::mutex mut_result_check;

    //Best weight found and the move that describes it
    int best_move_weight = std::numeric_limits<int>::min();
    int best_move_from = -1;
    int best_move_to = -1;
    bool best_move_makes_queen = false;

    //Initial alpha weight
    int a = std::numeric_limits<int>::min();

    //This function only does the maximizing part of minimax, its special because we need to know what the best move was

    Board board_copy = Board(board);

    //Go through the board and for each black piece, make a thread that will minimax
    for(int x = 0; x < BOARD_SIZE; x++)
    {
        for(int y = BOARD_SIZE-1; y >= 0; y--)
        {
            Piece cur_piece = board.get_piece(x,y);
            if (cur_piece.get_team() == BLACK)
            {
                threads.push_back(std::thread([&mut_result_check, &best_move_weight, &best_move_from, &best_move_to, &best_move_makes_queen, &a, x, y, &board_copy](){
                    //Get the possible moves of this piece
                    std::vector<int> piece_moves = board_copy.get_moves_for_space(x,y, true);
                    Piece cur_piece = board_copy.get_piece(x,y);
                    for (unsigned int i = 0; i < piece_moves.size(); i++)
                    {
                        //Make a test board and do the move on it
                        Board test_board(board_copy);
                        test_board.do_move(x+y*BOARD_SIZE, piece_moves[i], false);

                        //Minimax on this test board
                        int board_value = minimax(test_board, 4, 5, a, std::numeric_limits<int>::max(), false);

                        //Record new bests if they are found
                        mut_result_check.lock();
                        if (board_value > best_move_weight)
                        {
                            best_move_weight = board_value;
                            best_move_from = x+y*BOARD_SIZE;
                            best_move_to = piece_moves[i];
                            best_move_makes_queen = false;
                        }

                        //Update alpha for further iterations
                        if (best_move_weight > a)
                        {
                            a = best_move_weight;
                        }
                        mut_result_check.unlock();
                    }
                }));
            }
        }
    }

    //Wait for all threads to finish
    for(std::thread& cur_thread : threads)
    {
        cur_thread.join();
    }

    //Do the best found move
    board.do_move(best_move_from, best_move_to, best_move_makes_queen);

    return best_move_to;
}

//Regular alpha beta pruning, max depth is used to only make threads when <= 1 node from root
int Game::minimax(Board current_board, int depth, int max_depth, int a, int b, bool maximizing)
{
    if (depth == 0 || current_board.is_mate(current_board.get_active_color()))
    {
        return board_heuristic(current_board);
    }

    int best_move_weight = maximizing?std::numeric_limits<int>::min():std::numeric_limits<int>::max();
    bool a_eject = false, b_eject  = false;

    //Thread management
    std::vector<std::thread> threads;
    std::mutex mut_result_check;
    int finished_threads = 0;
    std::mutex mut_finished_threads;

    //Board iteration optimization
    int y_start=(maximizing)?(BOARD_SIZE-1):0, y_increment=(maximizing)?-1:1;

    //Loop board, find pieces and minimax on their moves, make threads if available
    for(int x = 0; x < BOARD_SIZE && !a_eject && !b_eject; x++)
    {
        for(int y = y_start; y >= 0 && y < BOARD_SIZE && !a_eject && !b_eject; y+=y_increment)
        {
            if (current_board.get_piece(x,y).get_team() == (maximizing?BLACK:WHITE))
            {
                //Try to make a thread?
                if (depth == max_depth-1)
                {
                    //Is a thread available?
                    if (threads.size()-finished_threads <= 4)
                    {
                        threads.push_back(std::thread([&finished_threads, &mut_finished_threads, &mut_result_check, current_board, x, y, depth, max_depth, &best_move_weight, &a, &b, &a_eject, &b_eject, maximizing](){
                            minimax_evaluate(mut_result_check, current_board, x, y, depth, max_depth, best_move_weight, a, b, a_eject, b_eject, maximizing);
                            mut_finished_threads.lock();
                            finished_threads++;
                            mut_finished_threads.unlock();
                        }));
                    }
                    else
                    {   //Run in this thread
                        minimax_evaluate(mut_result_check, current_board, x, y, depth, max_depth, best_move_weight, a, b, a_eject, b_eject, maximizing);
                    }
                }
                else //Cannot make threads, run in this thread
                {
                    minimax_evaluate(mut_result_check, current_board, x, y, depth, max_depth, best_move_weight, a, b, a_eject, b_eject, maximizing);
                }
            }
        }
    }

    //Wait for threads to finish
    for (std::thread& cur_thread : threads)
    {
        cur_thread.join();
    }

    //Return best result
    return best_move_weight;
}

//Evaluation of each pieces moves through minimax
void Game::minimax_evaluate(std::mutex& mut_result_check, Board current_board, int x, int y, int depth, int max_depth, int& best_move_weight, int& a, int& b, bool& a_eject, bool& b_eject, bool maximizing)
{   
    Piece cur_piece = current_board.get_piece(x,y);
    std::vector<int> piece_moves = current_board.get_moves_for_space(x,y, true);

    if (maximizing)
    {
        for (unsigned int i = 0; i < piece_moves.size(); i++)
        {
            //Go through each move, make a test board and maximize off of it

            Board test_board(current_board);
            test_board.do_move(x+y*BOARD_SIZE, piece_moves[i], false);

            int board_value = minimax(test_board, depth-1, max_depth, a, b, false);

            mut_result_check.lock();
            if (board_value > best_move_weight)
            {
                best_move_weight = board_value;
            }

            if (best_move_weight >= b)
            {
                b_eject = true;
                mut_result_check.unlock();
                return;
            }

            if (best_move_weight > a)
            {
                a = best_move_weight;
            }
            mut_result_check.unlock();
        }
    }
    else
    {
        for (unsigned int i = 0; i < piece_moves.size(); i++)
        {
            //Go through each move, make a test board and minimize off of it

            Board test_board(current_board);
            test_board.do_move(x+y*BOARD_SIZE, piece_moves[i], false);

            int board_value = minimax(test_board, depth-1, max_depth, a, b, true);

            mut_result_check.lock();
            if (board_value < best_move_weight)
            {
                best_move_weight = board_value;
            }

            if (best_move_weight <= a)
            {
                a_eject = true;
                mut_result_check.unlock();
                return;
            }

            if (best_move_weight < b)
            {
                b = best_move_weight;
            }
            mut_result_check.unlock();
        }
    }
}

int Game::board_heuristic(Board current_board)
{
    int heuristic = 0;
    int piece_weights[2][7] = {{-900, -90, -30, -30, -50, -10, 0}, {900, 90, 30, 30, 50, 10, 0}};
    
    //Unknown piece weight calculation
    int unknown_white_count = 0;
    int unknown_black_count = 0;
    for(int x = 0; x < BOARD_SIZE; x++)
    {
        for(int y = 0; y < BOARD_SIZE; y++)
        {
            Piece cur_piece = current_board.get_piece(x,y);

            if (cur_piece.get_type() == NO_TYPE)
            {
                continue;
            }

            if (cur_piece.get_vis() == HIDDEN)
            {
                if (cur_piece.get_team() == WHITE)
                {
                    piece_weights[WHITE][UNKNOWN] += 2*piece_weights[WHITE][cur_piece.get_type()];
                    unknown_white_count++;
                }
                else
                {
                    piece_weights[BLACK][UNKNOWN] += 2*piece_weights[BLACK][cur_piece.get_type()];
                    unknown_black_count++;
                }
            }
        }
    }

    if (unknown_white_count != 0)
    {
        piece_weights[WHITE][UNKNOWN] /= unknown_white_count;
    }
    if (unknown_black_count != 0)
    {
        piece_weights[BLACK][UNKNOWN] /= unknown_black_count;
    }

    //Per piece weight calculation
    for(int x = 0; x < BOARD_SIZE; x++)
    {
        for(int y = 0; y < BOARD_SIZE; y++)
        {
            Piece cur_piece = current_board.get_piece(x,y);

            if (cur_piece.get_type() == NO_TYPE)
            {
                continue;
            }

            //Mark hidden pieces as unknown and not their actual type
            if (cur_piece.get_vis() == HIDDEN)
            {
                cur_piece = Piece(UNKNOWN, cur_piece.get_team(), SHOWN);
            }

            //Get unmodified weight
            int adden = piece_weights[cur_piece.get_team()][cur_piece.get_type()];

            //Pawn rank bonus
            if (cur_piece.get_type() == PAWN)
            {
                if (cur_piece.get_team() == WHITE)
                {
                    adden -= 2*(6-y);
                }
                else
                {
                    adden += 2*(y-1);
                }
            }

            //Attackers and defenders bonus
            int found_attackers = current_board.get_attackers_for_space(x+y*BOARD_SIZE);

            if (found_attackers > 0)
            {
                adden /= 2;
            }
            else if (found_attackers < 0)
            {
                adden = 4*adden/3;
            }

            heuristic += adden;
        }
    }

    return heuristic;
}