#include <iostream>
#include <vector>
#include <algorithm>
#include <thread>
#include <atomic>
#include <chrono>
#include <type_traits>

#include "types.h"
#include "Board.h"
#include "Move.h"
#include "uci.h"
#include "syncstream.h"
#include "search.h"
#include "HashTable.h"

namespace Search {

// Node of a transposition table and its type
enum TTType { exact, upper_bound, lower_bound, unknown };
struct TTEntry {
    enum TTType type;
    unsigned age;
    int depth;
    Score score;
    Move move;
};
// Transposition Table collision policy
struct TTPolicy {
    bool operator()(const TTEntry& entry, const TTEntry& candidate) {
        return (entry.age < candidate.age)
            || ((entry.type == candidate.type) && (entry.depth <= candidate.depth))
            || (candidate.type == exact);
    }
};

// Implementation specific global search state, workers and concurrency controls
State state;
HashTable<Board, TTEntry, TTPolicy> tt;
std::atomic<bool> isRunning;
std::vector<std::thread> workers;

// Set initial static age of search
template <typename Evaluator>
unsigned PVS<Evaluator>::_age = 0;

void init(unsigned ttSize) {
    // Init internal state
    state.depth = 16;                   // default max search depth
    state.wtime = 24 * 60 * 60 * 1000;  // 1 day in milliseconds
    state.btime = 24 * 60 * 60 * 1000;  // 1 day in milliseconds
    state.winc = 0;
    state.binc = 0;
    state.movetime = 0;

    // Init Transposition Table
    constexpr unsigned ttLineSize = sizeof(HashTable<Board, TTEntry, TTPolicy>::Line);
    // Reserve/resize tt to ttSize in MB of memory
    if (tt.reserve(ttSize * 1024U * 1024U / ttLineSize)) {
        // Report error of reserve failed
        osyncstream(cout) << "info string error Unable to resize "
            "Transposition Table to " << ttSize << "MB" << std::endl;
    }
}

double ttSize() {
    // Init Transposition Table
    constexpr unsigned ttLineSize = sizeof(HashTable<Board, TTEntry, TTPolicy>::Line);
    // compute size in MB from number of lines (size)
    return static_cast<double>(tt.size() * ttLineSize) / 1048576.0;
}

void newgame() {
    // Reset internal state
    state.depth = 16;                   // default max search depth
    state.wtime = 24 * 60 * 60 * 1000;  // 1 day in milliseconds
    state.btime = 24 * 60 * 60 * 1000;  // 1 day in milliseconds
    state.winc = 0;
    state.binc = 0;
    state.movetime = 0;

    // Delete all data from TT table
    tt.clear();
}

// Perft implementation
unsigned perft_subroutine(Board& board, int depth) {
    if (depth <= 0) { return 1; }

    Index nodeCount = 0;
    MoveList searchmoves;
    board.moves(searchmoves);

    if (depth == 1) {
        nodeCount = searchmoves.size();
    } else {
        Board boardCopy = board;
        for (Move move : searchmoves) {
            board.make(move);
            nodeCount += perft_subroutine(board, depth - 1);
            board = boardCopy; // unmake move
        }
    }

    return nodeCount;
}
/**
 * Prints for all moves in passed move list (assuming legal moves) number of
 * `depth - 1` performance test node count of this child node count.
 *
 * @param board Reference root node board position.
 * @param depth search depth.
 * @param moveList List of legal moves at current board position to search.
 *  If moveList is empty, search all legal moves (Note: In this case, the
 *  moveList is filled with all legal moves manipulating the list in place).
 *
 * @example Output generated from start position for depth `5` with moveList
 * containing the moves "d4", "Nf3" and "Nc3".
 * position startpos
 * go perft 5 searchmoves d2d4 g1f3 b1c3
 *      d2d4: 361790
 *      g1f3: 233491
 *      b1c3: 234656
 *
 *      Nodes searched: 829937
 */
void perft(Board& board, int depth, MoveList& searchmoves) {
    if (depth <= 0) {
        osyncstream(cout) << "Nodes searched: 0" << std::endl;
        return;
    }

    if (!searchmoves.size()) {
        board.moves(searchmoves);
    }

    Index totalCount = 0;

    Board boardCopy = board;
    for (Move move : searchmoves) {
        // check if shutdown/stopping is requested
        if (!isRunning.load(std::memory_order_relaxed)) { break; }
        // continue traversing
        board.make(move);
        Index nodeCount;
        nodeCount = perft_subroutine(board, depth - 1);
        totalCount += nodeCount;
        board = boardCopy; // unmake move

        // report moves of node
        osyncstream(cout) << move << ": " << nodeCount << std::endl;
    }

    osyncstream(cout)
        << std::endl << "Nodes searched: " << totalCount << std::endl;
}
/**
 * Specialization of `perft(Board&, int, MoveList&)` for all legal moves.
 */
void perft(Board& board, int depth) {
    MoveList moveList; // init empty move list -> perft for all legal moves
    perft(board, depth, moveList);
}

/**
 * Search initialization
 */
template <typename Evaluator>
PVS<Evaluator>::PVS(const std::vector<Board>& game, const State& config)
    : _root(game.back())
    // Restrict max search depth by the principal variation capacity
    // Upper depth bound allows more than 200 plys which is enough.
    , _max_depth{std::min(config.depth, static_cast<int>(MoveList::max_size()))}
    // Init selective depth (max reached pvs + qsearch depth, max search ply)
    , _seldepth{0}
    // Set visited nodes to zero
    , _nodes{0}
    // break condition flag
    , _isStopped{false}
    // search start time for time dependend break condition and reporting
    , _start_time{clock::now()}
    , _searchmoves(config.searchmoves)
{
    // Fill search moves with all legal moves if no moves are specified
    if (_searchmoves.empty()) {
        _root.moves(_searchmoves);
    }

    // Increment age (for TT entries)
    ++_age;

    // Time control; set time to stop searching (ignoring potental loss of 2 ms)
    unsigned half_move_time;
    if (config.movetime) {
        half_move_time = config.movetime / 2U;
    } else {
        if (_root.isWhiteTurn()) {
            half_move_time = std::max(config.wtime / 60U, 50U);
        } else {
            half_move_time = std::max(config.btime / 60U, 50U);
        }
    }
    _mid_time = _start_time + milliseconds(     half_move_time);
    _end_time = _start_time + milliseconds(2U * half_move_time);

    // fill repitition detection history
    _historySize = _root.halfMoveClock() + 1 < game.size()
        ? _root.halfMoveClock() + 1
        : game.size();
    auto gameIter = game.rbegin();
    for (size_t i = _historySize; i-- > 0; ) {
        _history[i] = (*gameIter++).hash();
    }
}

/**
 * Principal Variation Search at the Root Position
 */
template <typename Evaluator>
Score PVS<Evaluator>::operator()() {
    using std::chrono::duration_cast;

    // Color keeps track of static evaluation sign for the current player
    // in the negated maximization (negaMax) framework
    Score color = _root.isWhiteTurn() ? +1 : -1;

    // Tracks principal variation
    MoveList pv_line;

    // set killers all to !move (TODO: is this required?)
    std::fill_n(_killers[0], MoveList::max_size(), Move(0));
    std::fill_n(_killers[1], MoveList::max_size(), Move(0));

    // set history heuristic to zero    // TODO: can I replace that by zero-initialization?!
    for (Index p = 0; p < 6; ++p) {
        std::fill_n(_historyTable[0][p], 64, 0);
        std::fill_n(_historyTable[1][p], 64, 0);
    }

    // Iterative deepening
    Score prevScore = 0;
    for (int depth = 1; depth <= _max_depth; depth++) {
        // Clear temporary pv line
        pv_line.clear();

        // Aspiration window search
        Score alpha, score, beta;
        Score delta = 34;
        // Set initial aspriation window bounds (shallow searches -> full width)
        if (depth < 3) {
            alpha = limits<Score>::lower();
            beta  = limits<Score>::upper();
        } else {
            alpha = prevScore - delta;
            beta  = prevScore + delta;
        }
        while (true) {
            // search current depth and aspiration window
            score = pvs(_root, depth, 0, alpha, beta, color, pv_line);

            // Increase aspiration window
            delta += 2 * delta + 5;

            // Increase window (iff fail high/low)
            if (score <= alpha) {
                alpha = std::max(prevScore - delta, limits<Score>::lower() - 1);
            } else if (beta <= score) {
                beta = std::max(prevScore + delta, limits<Score>::upper() + 1);
            } else {
                break;
            }
        }

        // Getting deeper, set previous iteration score in iterative deepening
        prevScore = score;

        // Check search break condition (before updating the pv and reporting
        // partial results)
        if (_isStopped
        || !isRunning.load(std::memory_order_relaxed)
        || (_end_time < clock::now())) {
            break;
        }

        // Copy principal variation
        _pv.clear();
        _pv.append(pv_line);

        auto now = clock::now();
        // Time spend in search till now (only for reporting)
        auto duration = duration_cast<milliseconds>(now - _start_time);
        // Check if found mate and report mate instead of centi pawn score
        if (limits<Score>::upper() <= (std::abs(score) + depth)) {
            // Note: +-1 for converting 0-indexed ply count to a 1-indexed
            // move count
            if (score < 0) {
                score = limits<Score>::lower() - score - 1;
            } else {
                score = limits<Score>::upper() - score + 1;
            }
            // ply to move count
            score /= 2;
            // Report search stats and mate in score moves
            osyncstream(cout)
                << "info depth " << depth
                << " seldepth " << _seldepth
                << " score mate " << score
                << " time " << duration.count()
                << " nodes " << _nodes
                << " pv " << _pv
                << std::endl;
            // stop iterative deepening, found check mate -> no need to continue
            break;
        } else {
            // Report search progress
            osyncstream(cout)
                << "info depth " << depth
                << " seldepth " << _seldepth
                << " score cp " << score
                << " time " << duration.count()
                << " nodes " << _nodes
                << " pv " << _pv
                << std::endl;
        }

        // Check current time againt search mid time, iff passed the mid point
        // a deeper iteration in the iterative deepening will most likely not
        // finish. Therefore, stop and save the time for later.
        if (_mid_time <= now) { break; }
    }

    if (_pv.size()) {
        auto out = osyncstream(cout);
        out << "bestmove " << _pv[0];
        if (1 < _pv.size()) {
            out << " ponder " << _pv[1];
        }
        out << std::endl;
    } else {
        // If there is no move in _pv (can happen in extreme short time control)
        // just make a random (first) move, better than failing.
        _root.moves(_pv); // _pv is empty
        if (_pv.size()) {
            osyncstream(cout) << "bestmove " << _pv[0] << std::endl;
        } else {
            osyncstream(cout) << "bestmove !move" << std::endl;
        }
    }

    // Before leaving search, update search state (remaining time)
    auto duration = duration_cast<milliseconds>(clock::now() - _start_time);
    if (_root.isWhiteTurn()) {
        state.wtime += state.winc - (duration.count() + 1);
    } else {
        state.btime += state.binc - (duration.count() + 1);
    }

    // Reset running flag
    isRunning.store(false, std::memory_order_relaxed);

    return (_root.isWhiteTurn() ? 1 : -1) * prevScore;
}

template <class Evaluator>
Score PVS<Evaluator>::qSearch(const PVS<Evaluator>::Position& pos, int ply,
    Score alpha, Score beta, Score color
) {
    // Track selective search depth counter (reported while searching)
    _seldepth = std::max(_seldepth, ply);

    // Static evaluation
    Score score = color * pos.eval();

    bool capturesOnly = true;
    // Check if in check, if in check generate all moves (check evations).
    // Otherwise, check if not moving at all (stand pat) is already too good.
    if (pos.isCheck()) {
        capturesOnly = false;
    } else if (beta <= score) {
        ++_nodes;
        return beta;
    }

    // Generate all captures (captures only, if not in check)
    MoveList moveList;
    pos.moves(moveList, capturesOnly);

    // Check for terminal node
    if (moveList.empty()) {
        ++_nodes;
        // If there are no captures left, return static evaluation score
        if (capturesOnly) {
            return score;
        // When generated all moves (in case of check evations), the fact that
        // there are no moves means mate.
        } else {
            return color * (pos.isWhiteTurn() ? +1 : -1)
                         * (limits<Score>::lower() + ply);
        }
    }

    // Ensure capture is better than stand pat
    alpha = std::max(alpha, score);

    // if (capturesOnly) {
    //     for (Move& move : moveList) {
    //         // add 65538 to ensure the set move score is positive
    //         move.setScore(pos.see(move) + 65538);
    //     }
    // } else {
    for (Move& move : moveList) {
        // MVV-LVA capture move order
        move.setScore(move.calcScore());
    }
    // }
    std::sort(moveList.begin(), moveList.end());

    // Recursively capture pieces
    Position board = pos;
    for (Move move : moveList) {
        // SEE pruning (drop losing captures)
        if (capturesOnly && (pos.see(move) < 0)) {
            continue;
        }

        board.make(move);

        score = -qSearch(board, ply + 1, -beta, -alpha, -color);

        if (beta <= score) {
            return beta;    // fail-hard beta-cutoff
        }
        alpha = std::max(alpha, score);

        board = pos; // unmake
    }

    return alpha;
}

/**
 * Principal Variation Search routine (all non-root) positions
 */
template <class Evaluator>
Score PVS<Evaluator>::pvs(const PVS<Evaluator>::Position& pos, int depth, int ply,
    Score alpha, Score beta, Score color, MoveList& pv_line
) {
    // Check for search shutdown and in case of a repetition return draw score
    if (_isStopped || isRepetition(pos)) { return 0; }
    // Check break condition every few thousend nodes
    if (!(_nodes % (32 * 1024))) {
        // Just return 0, partial search will be discarded
        if (!isRunning.load(std::memory_order_relaxed)
        || (_end_time < clock::now())) {
            _isStopped = true;
            return 0;
        }
    }

    // Save original alpha bound before TT entry update
    Score original_alpha = alpha;

    // Move list containing all moves to be searched (usually all legal moves)
    MoveList moveList;
    if (!ply) {
        moveList.append(_searchmoves);
    }

    // Transposition table lookup
    auto tt_line = tt.find(pos);
    if ((tt_line != tt.end()) && ((*tt_line).depth >= depth)) {
        // Need to create all moves here to check if the TT move is legal
        // (or restricted to searchmoves)
        if (moveList.empty()) {
            pos.moves(moveList);
        }
        // Only use the TT entry if TT move is legal
        if (moveList.contains((*tt_line).move)) {
            switch ((*tt_line).type) {
                case lower_bound:
                    alpha = std::max(alpha, (*tt_line).score);
                    break;
                case exact:
                    ++_nodes;
                    pv_line.clear();
                    pv_line.push_back((*tt_line).move);
                    return (*tt_line).score;
                case upper_bound:
                    beta = std::min(beta, (*tt_line).score);
                    break;
                default:
                    assert((false && "Unknown TT node type entry in TT lookup"));
            }
            // Check altered alpha/beta bounds
            if (beta <= alpha) {
                ++_nodes;
                return (*tt_line).score;
            }
        } else {
            // Reset TT entry to "not-found"
            tt_line = tt.end();
        }
    }

    // Max search depth reached
    if (depth <= 0) {
        return qSearch(pos, ply, alpha, beta, color);
    }

    // Generate all moves (if not already generated or constraint)
    if (moveList.empty()) {
        pos.moves(moveList);
    }

    // Check terminal node
    if (moveList.empty()) {
        // Increment node count
        ++_nodes;
        // Evaluate terminal node score (check/stale mate)
        if (pos.isCheck()) {
            return color * (pos.isWhiteTurn() ? +1 : -1)
                         * (limits<Score>::lower() + ply);
        } else {
            return 0;
        }
    }

    // Sort moves, iff hash move its the first, then captures in MVV-LVA
    // followed by the rest
    auto moveOrder = moveList.begin();
    if (tt_line != tt.end()) {
        if (moveList.move_front((*tt_line).move)) {
            ++moveOrder;
        }
    }

    // score moves
    for (Move& move : moveList) {
        auto mScore = move.calcScore();
        // Add killer move score -> killers after captures
        mScore += (move == _killers[0][ply]) * Move::killerScore[0];
        mScore += (move == _killers[1][ply]) * Move::killerScore[1];
        // add history heuristic score -> sort of non-captures
        mScore += static_cast<bool>(move.victim())
                * _historyTable[move.color()][move.piece() - 2][move.to()];
        move.setScore(mScore);
    }

    // TODO: best move selection instead of sort
    std::sort(moveOrder, moveList.end(),
        [](Move::base_type lhs, Move::base_type rhs) { return lhs > rhs; });


    // New child pv line
    MoveList cur_line;

    // New TT entry (tracks "best" move)
    TTEntry tt_entry{unknown, _age, depth, 0, Move{}};

    // Current board
    Position board = pos;

    Index moveCount = 0;
    Score score = limits<Score>::lower();
    for (Move move : moveList) {
        board.make(move);
        pushHistory(board);

        Score value;
        if (!moveCount) {
            value = -pvs(board, depth - 1, ply + 1,
                         -beta, -alpha, -color, cur_line);
        } else {
            // LMR (Late Move Reduction), reduce search depth of later moves
            int R = 0;  // TODO: disabled!!!
            // int R = (moveCount > 3) & (ply > 2) & !move.victim() & !pos.isCheck();

            // Zero-window search (with possible reduction)
            value = -pvs(board, depth - 1 - R, ply + 1,
                         -(alpha + 1), -alpha, -color, cur_line);
            if (alpha < value) {
                // re-search (with clean current pv line and without reduction)
                cur_line.clear();
                value = -pvs(board, depth - 1, ply + 1,
                             -beta, -alpha, -color, cur_line);
            }
        }
        ++moveCount;
        score = std::max(score, value);

        // Check for beta cutoff
        if (beta <= score) {
            // Set possible lower bound hash hash move
            tt_entry.move = move;
            // Store killer move which is a quiet move causing a beta cut-off
            if (!move.victim()) {
                // If move isn't the first killer move we insert which ensures
                // (assuming killers are different or !move) that two different
                // killers are stored
                if (_killers[0][ply] != move) {
                    _killers[1][ply] = _killers[0][ply];
                    _killers[0][ply] = move;
                }
                // Increment history heuristic
                _historyTable[move.color()][move.piece() - 2][move.to()] += depth * depth;
            }
            popHistory();
            break;
        }
        // Check if alpha raised; iff add new pv/best move
        if (alpha < score) {
            alpha = score;
            // Track pv (implicitly the best move)
            pv_line.clear();
            pv_line.push_back(move);
            pv_line.append(cur_line);
            // Track best move as hash move
            tt_entry.move = move;
        }

        board = pos; // unmake move
        popHistory();
    }

    // Add TT entry
    tt_entry.score = score;
    if (score <= original_alpha) {
        tt_entry.type = upper_bound;
    } else if (beta <= score) {
        tt_entry.type = lower_bound;
    } else {
        tt_entry.type = exact;
    }
    tt.insert(pos, tt_entry);

    return score;
}


// Starts searching the given board position by dispatching a worker thread
void start(std::vector<Board>& game, State& config) {
    if (isRunning.load(std::memory_order_consume)) {
        // Ignore further search start attempts if allready searching, requires
        // a stop first!
        return;
    } else {
        // Dispose of finished/stopped workers
        for (std::thread& worker : workers) {
            if (worker.joinable()) {
                worker.join();
            }
        }
        workers.clear();
    }
    // sets worker thread stop condition to false (before dispatch)
    isRunning.store(true, std::memory_order_release);

    // Dispatch search worker
    switch (config.search) {
        case State::Type::perft:
            state.search = State::Type::perft;
            workers.emplace_back([&]() {
                // Copy working variables, not subject to change
                Board pos(game.back());
                MoveList searchmoves = config.searchmoves;
                // Start perft
                perft(pos, config.depth, searchmoves);
                // Reset running flag
                isRunning.store(false, std::memory_order_relaxed);
            });
            break;
        case State::Type::search:
            std::cout << "info string HCE evaluation" << std::endl;
            workers.emplace_back(PVS<Board>(game, config));
            break;
        case State::Type::ponder:
            std::cerr << "info string error pondering not implemented!" << std::endl;
            break;
        default:
            assert((false && "Search::start got request for unknown search type."));
    }
}

void stop() {
    // revoke isRunning flag -> workers stop as soon as possible
    isRunning.store(false, std::memory_order_relaxed);

    // then join and dispose all workers
    for (std::thread& worker : workers) {
        if (worker.joinable()) {
            worker.join();
        }
    }
    workers.clear();
}

// Explicit instantiations of PVS types using different evaluators
template class PVS<Board>;

} /* namespace Search */
