bonsai.cpp

#include <stdio.h>
#include <iostream>
#include <math.h>
#include <iomanip>
#include <bit>
#include <fstream>
#include <chrono>
#include <algorithm>
#include <queue>
#include <immintrin.h>
#include <random>
#include "include/chess.hpp"

using namespace chess;

struct Node
{
    Move action;
    int firstchild; int num_children;
    int visits = 0;
    float value = 0;
    float policy;
};
struct SearchResult {
    Move bestMove;
    float confidence;

    int nps;

    // policy target
    std::vector<std::pair<Move, int>> policy;
};

Board board;
std::vector<Node> tree;

bool debug = false;
bool datagenActive = false;

const int DEFAULT_HASH_MIB = 128;
int hashMib = DEFAULT_HASH_MIB;
size_t hashLimitNodes = static_cast<size_t>(DEFAULT_HASH_MIB) * 1024 * 1024 / sizeof(Node);

inline size_t computeHashLimitNodes(int mib) {
    return static_cast<size_t>(mib) * 1024 * 1024 / sizeof(Node);
}

// Fraction of total nodes pruneTree() retains per call (~half).
const float PRUNE_KEEP_FRACTION = 0.5f;

#include "net/iteration1.hpp"

float sigmoid(float x) {
    return 1.0f / (1.0f + std::exp(-x));
}

inline float evaluate_network_16hl(const chess::Board& b) {
    int32_t acc[16];
    for (int i = 0; i < 16; ++i) {
        acc[i] = HIDDEN_BIASES[i];
    }

    chess::Color stm = b.sideToMove();
    bool isBlack = (stm == chess::Color::BLACK);

    // Precompute rank/file lookup to avoid constructing chess::Square twice
    for (int sq = 0; sq < 64; ++sq) {
        chess::Piece piece = b.at(chess::Square(sq));
        if (piece == chess::Piece::NONE) continue;

        int type_idx;
        switch (piece.type()) {
            case PieceType(PieceType::PAWN):   type_idx = 0; break;
            case PieceType(PieceType::KNIGHT): type_idx = 1; break;
            case PieceType(PieceType::BISHOP): type_idx = 2; break;
            case PieceType(PieceType::ROOK):   type_idx = 3; break;
            case PieceType(PieceType::QUEEN):  type_idx = 4; break;
            case PieceType(PieceType::KING):   type_idx = 5; break;
            default: continue;
        }

        int piece_idx = type_idx;
        if (piece.color() != stm) {
            piece_idx += 6;
        }

        // sq is 0-63: rank = sq / 8, file = sq % 8  (algebraic: a1=0, h1=7, a8=56, h8=63)
        int rank = sq >> 3;      // sq / 8
        int file = sq & 7;       // sq % 8
        int python_square = (7 - rank) * 8 + file;
        if (isBlack) {
            python_square ^= 56;
        }

        int feature_idx = piece_idx * 64 + python_square;

        // Accumulate contribution of this single active feature into all 16 hidden neurons
        for (int h = 0; h < 16; ++h) {
            acc[h] += HIDDEN_WEIGHTS[h][feature_idx];
        }
    }

    // SCReLU
    constexpr int32_t CLAMP_LIMIT = 255; 
    int32_t output_accumulation = OUTPUT_BIAS;

    for (int i = 0; i < 16; ++i) {
        int32_t h = acc[i];
        
        // Branchless clamp to [0, CLAMP_LIMIT]
        h = (h < 0) ? 0 : ((h > CLAMP_LIMIT) ? CLAMP_LIMIT : h);
        
        // Square the clipped value and accumulate into output
        output_accumulation += h * h * OUTPUT_WEIGHTS[i];
    }

    float raw_output_float = static_cast<float>(output_accumulation) / SCALE_FACTOR_SQ;
    return sigmoid(raw_output_float);
}

// obsolete?
inline bool isNodeTerminal(GameResult result) {
    return result != GameResult::NONE;
}

bool isNodeFullyExpanded(int node) {
    return tree[node].firstchild != 0;
}

void expandNode(int parent) {
    Movelist moves;
    movegen::legalmoves(moves, board);

    if (moves.empty()) return;

    constexpr int pieceValues[] = {100, 300, 300, 500, 900, 10000};

    auto getMvvLvaScore = [&](const Move& m) -> int {
        bool isEnPassant = (m.typeOf() == Move::ENPASSANT);
        Piece victimPiece = board.at(m.to());

        if (victimPiece == Piece::NONE && !isEnPassant) {
            return 0;
        }

        PieceType attacker = board.at(m.from()).type();
        PieceType victim = isEnPassant ? PieceType::PAWN : victimPiece.type();

        return (pieceValues[static_cast<int>(victim)] * 10) 
               - pieceValues[static_cast<int>(attacker)] 
               + 100000;
    };

    std::sort(moves.begin(), moves.end(), [&](const Move& a, const Move& b) {
        return getMvvLvaScore(a) > getMvvLvaScore(b);
    });

    tree[parent].firstchild = tree.size();
    tree[parent].num_children = moves.size();

    std::vector<double> policyWeights(moves.size());
    double totalWeight = 0.0;

    for (size_t i = 0; i < moves.size(); ++i) {
        policyWeights[i] = 1.0 / (static_cast<double>(i) + 1.0);
        totalWeight += policyWeights[i];
    }

    for (size_t i = 0; i < moves.size(); ++i) {
        Node child = Node();
        child.action = moves[i];
        child.policy = policyWeights[i] / totalWeight;
        tree.push_back(child);
    }
}

const float C_PUCT = 1.41f;

inline float fast_sqrt_hint(float x) {
    if (x <= 0.0f) return 0.0f;
    #if defined(__SSE2__)
    __m128 v = _mm_set_ss(x);
    __m128 r = _mm_rsqrt_ss(v); // r ~ 1/sqrt(x)
    
    // Newton-Raphson refinement for sqrt(x):
    // result = 0.5 * (x * r) * (3.0 - x * r * r)
    __m128 x_r = _mm_mul_ss(v, r);
    __m128 x_r_r = _mm_mul_ss(x_r, r);
    __m128 three = _mm_set_ss(3.0f);
    __m128 half = _mm_set_ss(0.5f);
    
    __m128 term = _mm_sub_ss(three, x_r_r);
    __m128 refined = _mm_mul_ss(half, _mm_mul_ss(x_r, term));
    
    float result;
    _mm_store_ss(&result, refined);
    return result;
    #else
    return sqrtf(x);
    #endif
}

// https://github.com/TomaszJaworski777/cpu-mcts-tutorial
inline float PUCT(int node, int parent) {
    const Node& child = tree[node];
    const Node& parent_node = tree[parent];

    float q = (child.visits == 0) ? 0.5f : (child.value / child.visits);

    float N_parent = static_cast<float>(parent_node.visits);
    if (N_parent < 1.0f) N_parent = 1.0f;

    // Inject child.policy here to scale the exploration factor
    float p = child.policy * fast_sqrt_hint(N_parent) / (static_cast<float>(child.visits) + 1.0f);

    return q + C_PUCT * p;
}

int selectBestChild(int parent) {
    int bestChild = tree[parent].firstchild;
    float bestChildPUCT = -MAXFLOAT; // no MINFLOAT?

    for(int i = 0; i < tree[parent].num_children; i++) {
        float currentChildPUCT = PUCT(tree[parent].firstchild + i, parent);

        if(currentChildPUCT > bestChildPUCT) {
            bestChild = tree[parent].firstchild + i;
            bestChildPUCT = currentChildPUCT;
        }
    }

    return bestChild;
}

void backpropagateResult(const std::vector<int>& line, float value) {
    // Backpropagate through the path from leaf to child of root
    for (int i = static_cast<int>(line.size()) - 1; i >= 0; i--) {
        tree[line[i]].visits++;
        tree[line[i]].value += value;
        value = 1 - value;
    }
    // Update root node
    tree[0].visits++;
    tree[0].value += value;
}

// theoretically unused but leaving it here for static evaluation
inline int material(const chess::Board& b) {
    auto count = [](auto bb) { return std::popcount(bb.getBits()); };   

    return
        100 * (count(b.us(chess::Color::WHITE) & b.pieces(chess::PieceType::PAWN))
             - count(b.us(chess::Color::BLACK) & b.pieces(chess::PieceType::PAWN))) +
        320 * (count(b.us(chess::Color::WHITE) & b.pieces(chess::PieceType::KNIGHT))
             - count(b.us(chess::Color::BLACK) & b.pieces(chess::PieceType::KNIGHT))) +
        330 * (count(b.us(chess::Color::WHITE) & b.pieces(chess::PieceType::BISHOP))
             - count(b.us(chess::Color::BLACK) & b.pieces(chess::PieceType::BISHOP))) +
        500 * (count(b.us(chess::Color::WHITE) & b.pieces(chess::PieceType::ROOK))
             - count(b.us(chess::Color::BLACK) & b.pieces(chess::PieceType::ROOK))) +
        900 * (count(b.us(chess::Color::WHITE) & b.pieces(chess::PieceType::QUEEN))
             - count(b.us(chess::Color::BLACK) & b.pieces(chess::PieceType::QUEEN)));
}

// Helper: build PV line (most-visited child path) from root into `out`
void buildPV(std::vector<int>& out) {
    out.clear();
    int cur = 0;
    while (isNodeFullyExpanded(cur)) {
        int bestChild = 0;
        int bestVisits = -1;
        for (int i = 0; i < tree[cur].num_children; i++) {
            int idx = tree[cur].firstchild + i;
            if (tree[idx].visits > bestVisits) {
                bestChild = idx;
                bestVisits = tree[idx].visits;
            }
        }
        out.push_back(bestChild);
        cur = bestChild;
    }
}

// Helper: print UCI info line
void printInfo(long long plyTotal, int plyMax, int iterationsCompleted,
               double safeElapsed, const std::vector<int>& pvLine) {
    int simPerSec = static_cast<int>(iterationsCompleted / safeElapsed);
    int hashfull = 0;
    if (hashLimitNodes > 0) {
        size_t used = tree.size();
        if (used > hashLimitNodes) used = hashLimitNodes;
        hashfull = static_cast<int>((used * 1000) / hashLimitNodes);
    }
    float y = floor(((1.0f - (static_cast<float>(tree[0].value) / tree[0].visits)) - 0.5f) * 100.0f * 20.0f);

    std::cout << "info depth " << (plyTotal / iterationsCompleted)
              << " seldepth " << plyMax
              << " nodes " << iterationsCompleted
              << " nps " << simPerSec
              << " hashfull " << hashfull
              << " time " << static_cast<int>(safeElapsed * 1000)
              << " score cp " << y
              << " pv ";
    for (int node : pvLine) {
        std::cout << uci::moveToUci(tree[node].action) << " ";
    }
    std::cout << std::endl;
}

void pruneTree() {
    const size_t n = tree.size();
    if (n < 1024) return; // not worth compacting tiny trees

    size_t keepTarget = static_cast<size_t>(n * PRUNE_KEEP_FRACTION);
    if (keepTarget < 1) keepTarget = 1;
    if (keepTarget > n) keepTarget = n;

    // --- Pass A: mark the top-`keepTarget` nodes by (visits desc, index asc) ---
    // That set is ancestor-closed (see above), so it is a connected subtree
    // rooted at node 0 and compacts cleanly.
    std::vector<bool> kept(n, false);
    {
        std::vector<std::pair<int,int>> rank; // (visits, index)
        rank.reserve(n);
        for (size_t i = 0; i < n; i++)
            rank.push_back({tree[i].visits, static_cast<int>(i)});
        std::partial_sort(rank.begin(), rank.begin() + keepTarget, rank.end(),
                          [](const std::pair<int,int>& a, const std::pair<int,int>& b) {
                              if (a.first != b.first) return a.first > b.first;
                              return a.second < b.second;
                          });
        for (size_t i = 0; i < keepTarget; i++) kept[rank[i].second] = true;
        kept[0] = true; // root always retained (defensive)
    } // `rank` freed before compaction to keep peak memory down

    // --- Pass B: emit the kept subtree in BFS order ---
    // BFS lets us write each kept node's kept children as one contiguous block
    // (same layout expandNode() produces), so firstchild + num_children stay
    // valid. Subtrees may land anywhere in the array -- that's fine, they are
    // reached through firstchild, exactly like the organically grown tree.
    size_t keptCount = 0;
    for (size_t i = 0; i < n; i++) if (kept[i]) keptCount++;

    std::vector<Node> newTree;
    newTree.reserve(keptCount);
    std::vector<int> oldToNew(n, -1);
    std::vector<int> keptKids; // reused
    std::queue<int> bfs;

    oldToNew[0] = 0;
    newTree.push_back(tree[0]);
    bfs.push(0);

    while (!bfs.empty()) {
        int p = bfs.front(); bfs.pop();
        int pNew = oldToNew[p];
        const Node& pnode = tree[p];

        keptKids.clear();
        if (pnode.firstchild != 0) {
            for (int i = 0; i < pnode.num_children; i++) {
                int ci = pnode.firstchild + i;
                if (kept[ci]) keptKids.push_back(ci);
            }
        }

        if (keptKids.empty()) {
            // No retained children: this node becomes a re-expandable leaf.
            newTree[pNew].firstchild = 0;
            newTree[pNew].num_children = 0;
            continue;
        }

        int childBlockStart = static_cast<int>(newTree.size());
        newTree[pNew].firstchild = childBlockStart;
        newTree[pNew].num_children = static_cast<int>(keptKids.size());
        for (int ci : keptKids) {
            oldToNew[ci] = static_cast<int>(newTree.size());
            newTree.push_back(tree[ci]); // firstchild fixed when ci is processed
            bfs.push(ci);
        }
    }

    if (debug && !datagenActive) {
        std::cout << "info string prune " << n << " -> " << newTree.size()
                  << " nodes (hash limit " << hashLimitNodes << ")" << std::endl;
    }

    tree = std::move(newTree);
}

SearchResult monteCarloSearch(int iterationsMax, int timeMax) {
    if (iterationsMax == 0) iterationsMax = INT32_MAX;

    tree.clear();

    Node root = Node();
    tree.push_back(root);
    expandNode(0);

    auto startTime = std::chrono::steady_clock::now();

    long long plyCurrent = 0;
    long long plyTotal = 0;
    int plyMax = 0;
    int iterationsCompleted = 0;

    // Reusable buffers to avoid per-iteration allocations
    std::vector<int> line;
    line.reserve(64);
    std::vector<int> pvLine;
    pvLine.reserve(32);

    int infoInterval = 100000;
    int nextInfo = infoInterval;
    int nextTimeCheck = 1000;

    for (; iterationsCompleted < iterationsMax; iterationsCompleted++) {
        line.clear();
        int current = 0;

        // Selection: descend to a leaf
        while (isNodeFullyExpanded(current)) {
            current = selectBestChild(current);
            board.makeMove(tree[current].action);
            plyCurrent++;
            line.push_back(current);
        }

        // Expansion + one more selection step
        auto [over, r] = board.isGameOver();
        if (r == GameResult::NONE && !isNodeFullyExpanded(current)) {
            expandNode(current);
            current = selectBestChild(current);
            board.makeMove(tree[current].action);
            plyCurrent++;
            line.push_back(current);
            // Re-check game state after the extra move
            auto result2 = board.isGameOver();
            r = result2.second;
        }

        // Track depth stats
        if (plyCurrent > plyMax) plyMax = plyCurrent;
        plyTotal += plyCurrent;
        plyCurrent = 0;

        // Evaluation
        float value = 0;
        switch (r) {
            case GameResult::WIN:  value = 1.0f; break;
            case GameResult::DRAW: value = 0.49f; break; // contempt
            case GameResult::LOSE: value = 0.0f; break;

            // removed * 0.9f, either tune or throw out, is probably bad for eval
            //case GameResult::NONE: value = evaluate_network_16hl(board); break;

            // fallback
            case GameResult::NONE: {
                float white_persp = 1.0 / (1.0 + std::exp(-static_cast<double>(material(board)) / 400.0));
                value = (board.sideToMove() == Color::WHITE) ? white_persp : (1.0f - white_persp);
                break;
            }
        }

        backpropagateResult(line, 1.0f - value);

        // Unmake moves to restore root position
        for (int i = static_cast<int>(line.size()) - 1; i >= 0; i--) {
            board.unmakeMove(tree[line[i]].action);
        }

        // Periodic time check (not every iteration)
        if (timeMax != 0 && iterationsCompleted >= nextTimeCheck) {
            auto currentTime = std::chrono::steady_clock::now();
            double elapsed = std::max(1e-9, std::chrono::duration<double>(currentTime - startTime).count());
            if (timeMax < elapsed * 1000.0) break;
            nextTimeCheck += 1000;
        }

        // compact
        if (tree.size() >= hashLimitNodes) {
            pruneTree();
        }

        // Periodic info print
        if (!datagenActive && iterationsCompleted >= nextInfo) {
            auto currentTime = std::chrono::steady_clock::now();
            double safeElapsed = std::max(1e-9, std::chrono::duration<double>(currentTime - startTime).count());
            buildPV(pvLine);
            printInfo(plyTotal, plyMax, iterationsCompleted, safeElapsed, pvLine);
            nextInfo += infoInterval;
        }
    }

    // Final PV and info
    buildPV(pvLine);

    auto currentTime = std::chrono::steady_clock::now();
    double safeElapsed = std::max(1e-9, std::chrono::duration<double>(currentTime - startTime).count());

    if (!datagenActive) {
        printInfo(plyTotal, plyMax, iterationsCompleted, safeElapsed, pvLine);
    }

    // Debug output
    if (debug && !datagenActive) {
        const int BAR_WIDTH = 40;
        std::cout << "\nMove Statistics:\n";
        std::cout << "-------------------------------\n";

        for (int i = 0; i < tree[0].num_children; i++) {
            const Node& child = tree[tree[0].firstchild + i];
            float ratio = static_cast<float>(child.visits) / iterationsCompleted;
            int filled = static_cast<int>(ratio * BAR_WIDTH);

            std::cout << " [";
            for (int j = 0; j < BAR_WIDTH; j++) {
                std::cout << (j < filled ? '#' : ' ');
            }
            std::cout << "] " << uci::moveToUci(child.action) << " "
                      << child.visits << " ("
                      << std::fixed << ratio * 100.0f << "%" << ") q="
                      << (child.visits > 0 ? child.value / child.visits : 0.0f)
                      << " puct=" << PUCT(tree[0].firstchild + i, 0) << "\n";
        }
    }

    int bestChild = pvLine.empty() ? tree[0].firstchild : pvLine[0];

    SearchResult result;
    result.bestMove = tree[bestChild].action;

    result.nps = static_cast<int>(iterationsCompleted / safeElapsed);
    
    result.confidence = 1.0f - (static_cast<float>(tree[0].value) / tree[0].visits);

    result.policy.reserve(tree[0].num_children);
    for (int i = 0; i < tree[0].num_children; i++) {
        const Node& child = tree[tree[0].firstchild + i];
        result.policy.push_back({child.action, child.visits});
    }

    return result;
}


struct DatagenPosition {
    std::string fen;
    std::vector<std::pair<Move, int>> policy;
    Color sideToMove;
    float confidence;
};

#pragma pack(push, 1)
struct PositionRecord {
    uint64_t bitboards[12]; // 96 bytes: 768-bit one-hot encoding
    float confidence;       // 4 bytes: MCTS search evaluation
};
#pragma pack(pop) // Total size: Exactly 100 bytes per position

// custom binary trickery

void writeGameToBinary(const std::string& filename,
                       const std::vector<DatagenPosition>& game,
                       const Color& winner)
{
    std::ofstream file(filename, std::ios::binary | std::ios::app);
    if (!file.is_open()) return;

    std::vector<PositionRecord> records;
    records.reserve(game.size());

    for (const auto& position : game)
    {
        PositionRecord record;

        float gameResult = 0.5f; // Default to a draw
        if (winner == Color::WHITE) {
            gameResult = (position.sideToMove == Color::WHITE) ? 1.0f : 0.0f;
        } else if (winner == Color::BLACK) {
            gameResult = (position.sideToMove == Color::BLACK) ? 1.0f : 0.0f;
        }
        
        // Apply lambda blend formula
        record.confidence = (0.7f * position.confidence) + (0.3f * gameResult);

        std::memset(record.bitboards, 0, sizeof(record.bitboards));

        bool isWhite = (position.sideToMove == Color::WHITE);
        int rank = 7;
        int fileIdx = 0;

        for (char c : position.fen)
        {
            if (c == ' ') break;
            
            if (c == '/') {
                rank--;
                fileIdx = 0;
            }
            else if (std::isdigit(c)) {
                fileIdx += c - '0';
            }
            else {
                int pIdx = -1;
                // Bake PIECE_MAP_W vs PIECE_MAP_B logic directly into the parser
                if (isWhite) {
                    switch (c) {
                        case 'P': pIdx = 0;  break; case 'N': pIdx = 1;  break;
                        case 'B': pIdx = 2;  break; case 'R': pIdx = 3;  break;
                        case 'Q': pIdx = 4;  break; case 'K': pIdx = 5;  break;
                        case 'p': pIdx = 6;  break; case 'n': pIdx = 7;  break;
                        case 'b': pIdx = 8;  break; case 'r': pIdx = 9;  break;
                        case 'q': pIdx = 10; break; case 'k': pIdx = 11; break;
                    }
                } else {
                    switch (c) {
                        case 'p': pIdx = 0;  break; case 'n': pIdx = 1;  break;
                        case 'b': pIdx = 2;  break; case 'r': pIdx = 3;  break;
                        case 'q': pIdx = 4;  break; case 'k': pIdx = 5;  break;
                        case 'P': pIdx = 6;  break; case 'N': pIdx = 7;  break;
                        case 'B': pIdx = 8;  break; case 'R': pIdx = 9;  break;
                        case 'Q': pIdx = 10; break; case 'K': pIdx = 11; break;
                    }
                }

                if (pIdx != -1) {
                    int square = rank * 8 + fileIdx;
                    // Flip the board vertically for Black's perspective (square ^ 56)
                    if (!isWhite) {
                        square ^= 56;
                    }
                    
                    record.bitboards[pIdx] |= (1ULL << square);
                    fileIdx++;
                }
            }
        }
        records.push_back(record);
    }

    file.write(reinterpret_cast<const char*>(records.data()), records.size() * sizeof(PositionRecord));
}

// viriformat
// spec: https://github.com/cosmobobak/viriformat

#pragma pack(push, 1)
struct ViriformatBoard {
    uint64_t occupied = 0;       // 8 bytes: Bitmask of occupied squares
    uint8_t pieces[16] = {0};    // 16 bytes: 32 packed 4-bit piece entries
    uint8_t stm_ep = 64;         // 1 byte: MSB = side to move, Lower 7 bits = EP square
    uint8_t halfmove = 0;        // 1 byte: Halfmove clock
    uint16_t fullmove = 1;       // 2 bytes: Fullmove counter
    int16_t score = 0;           // 2 bytes: Root position evaluation
    uint8_t result = 1;          // 1 byte: 0 = Black Win, 1 = Draw, 2 = White Win
    uint8_t padding = 0;         // 1 byte: Memory alignment padding
};
#pragma pack(pop) // Total size: Exactly 32 bytes

void writeGameToViriformat(const std::string& filename,
                           const std::vector<DatagenPosition>& game,
                           const Color& winner)
{
    if (game.empty()) return; // no moves happened

    // stream
    std::ofstream file(filename, std::ios::binary | std::ios::app);
    if (!file.is_open()) return;

    ViriformatBoard packed_board;
    
    // parse fen
    std::stringstream ss(game[0].fen);
    std::string pieces_str, stm_str, castling_str, ep_str, halfmove_str, fullmove_str;
    ss >> pieces_str >> stm_str >> castling_str >> ep_str >> halfmove_str >> fullmove_str;

    // Reconstruct initial position
    std::vector<char> board_state(64, ' ');
    int rank = 7;
    int fileIdx = 0;
    for (char c : pieces_str) {
        if (c == '/') {
            rank--;
            fileIdx = 0;
        } else if (std::isdigit(c)) {
            fileIdx += c - '0';
        } else {
            board_state[rank * 8 + fileIdx] = c;
            fileIdx++;
        }
    }

    // Set occupied bitboard
    for (int s = 0; s < 64; ++s) {
        if (board_state[s] != ' ') {
            packed_board.occupied |= (1ULL << s);
        }
    }

    // Pack piece entries into array
    int piece_count = 0;
    for (int s = 0; s < 64; ++s) {
        if (board_state[s] != ' ') {
            char c = board_state[s];
            bool is_black = std::islower(c);
            char type_char = std::tolower(c);
            
            uint8_t type = 0;
            if (type_char == 'p')      type = 0;
            else if (type_char == 'n') type = 1;
            else if (type_char == 'b') type = 2;
            else if (type_char == 'r') {
                type = 3; // Default Rook
                // Check castling strings to upgrade active rooks to Unmoved Rooks (Type 6)
                if (!is_black) {
                    if (s == 7 && castling_str.find('K') != std::string::npos) type = 6;
                    else if (s == 0 && castling_str.find('Q') != std::string::npos) type = 6;
                } else {
                    if (s == 63 && castling_str.find('k') != std::string::npos) type = 6;
                    else if (s == 56 && castling_str.find('q') != std::string::npos) type = 6;
                }
            }
            else if (type_char == 'q') type = 4;
            else if (type_char == 'k') type = 5;

            uint8_t piece_val = type;
            if (is_black) piece_val |= 0x08; // Set MSB bit for Black

            // Store two 4-bit entries per byte
            int byte_idx = piece_count / 2;
            if (piece_count % 2 == 0) {
                packed_board.pieces[byte_idx] |= (piece_val & 0x0F);
            } else {
                packed_board.pieces[byte_idx] |= ((piece_val & 0x0F) << 4);
            }
            piece_count++;
        }
    }

    // Pack Side-to-move and En-passant field
    uint8_t stm_ep = 64; 
    if (!ep_str.empty() && ep_str != "-") {
        stm_ep = (ep_str[1] - '1') * 8 + (ep_str[0] - 'a');
    }
    if (stm_str == "b") {
        stm_ep |= 0x80; // Set MSB if Black to move
    }
    packed_board.stm_ep = stm_ep;

    // Parsing clocks with safe defaults
    packed_board.halfmove = halfmove_str.empty() ? 0 : std::stoi(halfmove_str);
    packed_board.fullmove = fullmove_str.empty() ? 1 : std::stoi(fullmove_str);
    packed_board.score = 0; 

    // Game final result mapping
    if (winner == Color::WHITE)      packed_board.result = 2;
    else if (winner == Color::BLACK) packed_board.result = 0;
    else                             packed_board.result = 1;

    file.write(reinterpret_cast<const char*>(&packed_board), sizeof(ViriformatBoard));

    auto fill_flat_board = [](const std::string& fen, std::vector<char>& b) {
        std::stringstream fss(fen);
        std::string p_str; fss >> p_str;
        int r = 7, f = 0;
        for (char c : p_str) {
            if (c == '/') { r--; f = 0; }
            else if (std::isdigit(c)) { f += c - '0'; }
            else { b[r * 8 + f] = c; f++; }
        }
    };

    for (size_t i = 0; i < game.size() - 1; ++i) {
        std::vector<char> b1(64, ' '), b2(64, ' ');
        fill_flat_board(game[i].fen, b1);
        fill_flat_board(game[i + 1].fen, b2);

        bool is_white_turn = (game[i].sideToMove == Color::WHITE);
        int from_sq = -1, to_sq = -1;
        int move_type = 0;   // 0: Normal, 1: EP, 2: Castling, 3: Promotion
        int promo_piece = 0; // 0: N, 1: B, 2: R, 3: Q

        // king takes rook castling
        if (is_white_turn) {
            if (b1[4] == 'K' && b2[6] == 'K')      { from_sq = 4;  to_sq = 7;  move_type = 2; }
            else if (b1[4] == 'K' && b2[2] == 'K') { from_sq = 4;  to_sq = 0;  move_type = 2; }
        } else {
            if (b1[60] == 'k' && b2[62] == 'k')    { from_sq = 60; to_sq = 63; move_type = 2; }
            else if (b1[60] == 'k' && b2[58] == 'k') { from_sq = 60; to_sq = 56; move_type = 2; }
        }

        // If not a castling move, parse square deltas
        if (move_type == 0) {
            std::vector<int> changed;
            for (int s = 0; s < 64; ++s) {
                if (b1[s] != b2[s]) changed.push_back(s);
            }

            // Identify the active moving piece's starting square
            for (int s : changed) {
                if (b1[s] != ' ' && (std::isupper(b1[s]) == is_white_turn)) {
                    from_sq = s;
                }
            }

            if (changed.size() == 3) { 
                // En-Passant Capture
                move_type = 1;
                for (int s : changed) {
                    if (s != from_sq && b1[s] == ' ' && (b2[s] == 'P' || b2[s] == 'p')) {
                        to_sq = s;
                    }
                }
            } else {
                // Normal move, simple capture, or promotion
                for (int s : changed) {
                    if (s != from_sq) to_sq = s;
                }
                // Check for Pawn Promotion
                if (std::tolower(b1[from_sq]) == 'p' && (to_sq / 8 == 7 || to_sq / 8 == 0)) {
                    move_type = 3;
                    char p_char = std::tolower(b2[to_sq]);
                    if (p_char == 'n')      promo_piece = 0;
                    else if (p_char == 'b') promo_piece = 1;
                    else if (p_char == 'r') promo_piece = 2;
                    else if (p_char == 'q') promo_piece = 3;
                }
            }
        }

        // Pack values into the Move structure
        uint16_t packed_move = 0;
        packed_move |= (from_sq & 0x3F);
        packed_move |= ((to_sq & 0x3F) << 6);
        packed_move |= ((promo_piece & 0x03) << 12);
        packed_move |= ((move_type & 0x03) << 14);

        // Convert confidence float to viriformat White-Relative Centipawn Score (i16)
        float gameResVal = 0.5f; 
        if (winner == Color::WHITE)      gameResVal = (game[i].sideToMove == Color::WHITE) ? 1.0f : 0.0f;
        else if (winner == Color::BLACK) gameResVal = (game[i].sideToMove == Color::BLACK) ? 1.0f : 0.0f;

        float blended = (0.7f * game[i].confidence) + (0.3f * gameResVal);

        // Inverse sigmoid with scale 400 (reverse of 1/(1+exp(-x/400)))
        float clamped = std::max(1e-6f, std::min(1.0f - 1e-6f, blended));
        int16_t move_score = static_cast<int16_t>(400.0f * std::log(clamped / (1.0f - clamped)));

        file.write(reinterpret_cast<const char*>(&packed_move), sizeof(uint16_t));
        file.write(reinterpret_cast<const char*>(&move_score), sizeof(int16_t));
    }

    uint32_t terminator = 0;
    file.write(reinterpret_cast<const char*>(&terminator), sizeof(uint32_t));
}

void datagen() {
    datagenActive = true;
    int i = 0;
    long allPlies = 0;
    
    // Set up a proper random engine for weighted sampling
    std::random_device rd;
    std::mt19937 gen(rd());

    while (true) {
        // Make a random opening position
        board.setFen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");
        int numMoves = 5 + rand() % 5;

        for (int j = 0; j < numMoves; j++) {
            Movelist moves;
            movegen::legalmoves(moves, board);

            if (moves.empty()) {
                break;
            }

            int index = rand() % moves.size();
            board.makeMove(moves[index]);
        }
        std::string startingFen = board.getFen();

        std::vector<DatagenPosition> game;
        int ply = 0;

        // --- START TIMER ---
        auto startTime = std::chrono::high_resolution_clock::now();

        while (true) {
            Movelist moves;
            movegen::legalmoves(moves, board);

            if (board.isGameOver().second != GameResult::NONE)
                break;

            // Run the deep MCTS search for every recorded position
            SearchResult search = monteCarloSearch(1501, 0);

            DatagenPosition pos;
            pos.fen = board.getFen();
            pos.policy = search.policy;
            pos.sideToMove = board.sideToMove();
            pos.confidence = search.confidence;
            game.push_back(pos);

            Move moveToPlay;

            // --- TEMPERATURE MECHANISM ---
            if (ply < 20) { 
                // T = 1: Weighted random choice based strictly on visit counts
                int totalVisits = 0;
                for (const auto& p : search.policy) {
                    totalVisits += p.second;
                }

                if (totalVisits > 0) {
                    std::uniform_int_distribution<> dis(0, totalVisits - 1);
                    int target = dis(gen);
                    int currentSum = 0;

                    for (const auto& p : search.policy) {
                        currentSum += p.second;
                        if (currentSum > target) {
                            moveToPlay = p.first;
                            break;
                        }
                    }
                } else {
                    moveToPlay = search.bestMove; // Fallback safety
                }
            } 
            else { 
                // T = 0: Exploit mode. Play the absolute best move found by the search
                moveToPlay = search.bestMove;
            }

            board.makeMove(moveToPlay);
            ply++;
            allPlies++;
        }

        // --- END TIMER & CALCULATE SPEED ---
        auto endTime = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double> elapsed = endTime - startTime;

        auto [over, result] = board.isGameOver();

        Color winner = Color::NONE; // stays none if draw

        if (result == GameResult::LOSE)
        {
            // sideToMove is checkmated, so opposite side won
            winner = (board.sideToMove() == Color::WHITE)
                ? Color::BLACK
                : Color::WHITE;
        }
        
        double movesPerSec = (elapsed.count() > 0) ? (ply / elapsed.count()) : 0.0;

        std::cout << "Game " << i << ": " << startingFen 
                  << " | plies: " << ply 
                  << " | speed: " << std::fixed << std::setprecision(2) << movesPerSec << " plies/s" 
                  << " | total plies: " << allPlies
                  << std::endl;

        writeGameToViriformat("data/test.vf", game, winner);
        game.clear();
        i++;
    }
}

void handlePosition(std::istringstream& ss) {
    std::string token;
    ss >> token; // Should be "startpos" or "fen"

    if (token == "startpos") {
        board.setFen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");
    } else if (token == "fen") {
        std::string fen_string;
        // Read the FEN string- can have spaces
        while (ss >> token && token != "moves") {
            fen_string += token + " ";
        }
        if (!fen_string.empty()) {fen_string.pop_back();}
        board.setFen(fen_string);

        if (token == "moves") {
            while (ss >> token) {
                //std::cout << token << std::endl;
                chess::Move move = uci::uciToMove(board, token);
                board.makeMove(move);
            }
        }
        return;
    }

    ss >> token; // Should be "moves"
    if (token == "moves") {
        while (ss >> token) {
            //std::cout << token << std::endl;
            chess::Move move = uci::uciToMove(board, token);
            board.makeMove(move);
        }
    }
}

void benchAndQuit() {
    board.setFen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");
    SearchResult result = monteCarloSearch(100000, 0);
    int bench = result.policy[0].second + result.policy[1].second;
    std::cout << "Bench: " << bench << std::endl;
    std::cout << bench << " nodes " << result.nps << " nps" << std::endl;
    exit(0);
}

int main(int argc, char* argv[]) {
    std::string line;
    board = Board();
    board.setFen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");

    if(argc == 2) {
        benchAndQuit();
    }

    while (std::getline(std::cin, line)) {
        std::istringstream iss(line);
        std::string command;
        iss >> command;

        if (command == "uci") {
            std::cout << "id name Bonsai" << std::endl;
            std::cout << "id author Zander" << std::endl;
            std::cout << "option name Threads type spin default 1 min 1 max 1" << std::endl;
            std::cout << "option name Hash type spin default " << DEFAULT_HASH_MIB << " min 1 max 65536" << std::endl;
            std::cout << "uciok" << std::endl;
        } else if (command == "debug") {
            std::string arg;
            iss >> arg;
            if (arg == "on" || arg == "true" || arg == "1") {
                debug = true;
            } else if (arg == "off" || arg == "false" || arg == "0" || arg.empty()) {
                debug = false;
            }
        } else if (command == "isready") {
            std::cout << "readyok" << std::endl;
        } else if (command == "ucinewgame") {
            board.setFen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1");
        } else if (command == "setoption") {
            std::string token, name;
            iss >> token;   // "name"
            iss >> name;
            iss >> token;   // "value"
            if (name == "Hash") {
                iss >> hashMib;
                if (hashMib < 1) hashMib = 1;
                hashLimitNodes = computeHashLimitNodes(hashMib);
            }
        } else if (command == "position") {
            handlePosition(iss);
        } else if (command == "eval") {
            std::cout << "Evaluation for this position: " << material(board) << std::endl;
            //std::cout << "Evaluation for this position: " << evaluate_network(board) << std::endl;
        } else if (command == "go") {
            int nodes = 0;
            int movetime = 0; 
            int wtime = 0, btime = 0;
            int winc = 0, binc = 0;

            std::string token;
            while (iss >> token) {
                if (token == "nodes") {
                    iss >> nodes;
                } else if (token == "movetime") {
                    iss >> movetime;
                } else if (token == "wtime") {
                    iss >> wtime;
                } else if (token == "btime") {
                    iss >> btime;
                } else if (token == "winc") {
                    iss >> winc;
                } else if (token == "binc") {
                    iss >> binc;
                }
            }

            // Calculate time management if explicit movetime wasn't provided
            if (movetime == 0) {
                if (board.sideToMove() == Color::WHITE) {
                    if (wtime > 0) {
                        movetime = (wtime / 20) + winc;
                    }
                } else {
                    if (btime > 0) {
                        movetime = (btime / 20) + binc;
                    }
                }
            }

            Move bestMove = monteCarloSearch(nodes, movetime).bestMove;
            std::cout << "bestmove " << uci::moveToUci(bestMove) << std::endl;
        } else if (command == "dg") {
            datagen();
        } else if (command == "bench") {
            benchAndQuit();
        } else if (command == "quit") {
            break;
        }
    }

    return 0;
}