    // Copyright (C) 2025 Davis E. King (davis@dlib.net)
    // License: Boost Software License   See LICENSE.txt for the full license.
    #ifndef DLIB_BPE_TOKENIZER_H
    #define DLIB_BPE_TOKENIZER_H

    #include <iostream>
    #include <iomanip>
    #include <string>
    #include <vector>
    #include <future>
    #include <mutex>
    #include <thread>
    #include <map>
    #include <unordered_map>
    #include <queue>

    #include "../base64.h"
    #include "../serialize.h"
    #include "bpe_tokenizer_abstract.h"

    namespace dlib
    {
        constexpr size_t BPE_TOKENIZER_MAX_TOKEN_LENGTH = 8;
        constexpr int BPE_TOKENIZER_BASE_VOCAB_SIZE = 256;

        class bpe_tokenizer
        {
        public:
            bpe_tokenizer() : vocab_size(BPE_TOKENIZER_BASE_VOCAB_SIZE)
            {
                // Initialize the base vocabulary with single bytes
                for (int i = 0; i < BPE_TOKENIZER_BASE_VOCAB_SIZE; ++i)
                    vocab[i] = std::vector<uint8_t>{ static_cast<uint8_t>(i) };

                // Initialize special tokens with sequential IDs
                special_tokens =
                {
                    {"<text>",      BPE_TOKENIZER_BASE_VOCAB_SIZE},
                    {"</text>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 1},
                    {"<url>",       BPE_TOKENIZER_BASE_VOCAB_SIZE + 2},
                    {"</url>",      BPE_TOKENIZER_BASE_VOCAB_SIZE + 3},
                    {"<image>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 4},
                    {"</image>",    BPE_TOKENIZER_BASE_VOCAB_SIZE + 5},
                    {"<video>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 6},
                    {"</video>",    BPE_TOKENIZER_BASE_VOCAB_SIZE + 7},
                    {"<audio>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 8},
                    {"</audio>",    BPE_TOKENIZER_BASE_VOCAB_SIZE + 9},
                    {"<file>",      BPE_TOKENIZER_BASE_VOCAB_SIZE + 10},
                    {"</file>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 11},
                    {"<code>",      BPE_TOKENIZER_BASE_VOCAB_SIZE + 12},
                    {"</code>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 13},
                    {"<summary>",   BPE_TOKENIZER_BASE_VOCAB_SIZE + 14},
                    {"</summary>",  BPE_TOKENIZER_BASE_VOCAB_SIZE + 15},
                    {"<think>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 16},
                    {"</think>",    BPE_TOKENIZER_BASE_VOCAB_SIZE + 17},
                    {"<start>",     BPE_TOKENIZER_BASE_VOCAB_SIZE + 18},
                    {"<end>",       BPE_TOKENIZER_BASE_VOCAB_SIZE + 19},
                    {"<user>",      BPE_TOKENIZER_BASE_VOCAB_SIZE + 20},
                    {"<bot>",       BPE_TOKENIZER_BASE_VOCAB_SIZE + 21},
                    {"<system>",    BPE_TOKENIZER_BASE_VOCAB_SIZE + 22},
                    {"<question>",  BPE_TOKENIZER_BASE_VOCAB_SIZE + 23},
                    {"<answer>",    BPE_TOKENIZER_BASE_VOCAB_SIZE + 24},
                    {"<search>",    BPE_TOKENIZER_BASE_VOCAB_SIZE + 25},
                    {"<unk>",       BPE_TOKENIZER_BASE_VOCAB_SIZE + 26},
                    {"<pad>",       BPE_TOKENIZER_BASE_VOCAB_SIZE + 27}
                };

                // Initialize the vector of special token IDs
                for (const auto& token : special_tokens)
                    special_token_map[token.second] = token.first;
            }

            // Train the tokenizer on the given text
            void train(const std::string& text, int vocab_size, bool verbose = false)
            {
                DLIB_CASSERT(vocab_size >= BPE_TOKENIZER_BASE_VOCAB_SIZE);
                this->vocab_size = vocab_size;
                int num_merges = vocab_size - BPE_TOKENIZER_BASE_VOCAB_SIZE;

                // Convert text to byte IDs
                std::vector<int> ids;
                for (char c : text) ids.push_back(static_cast<uint8_t>(c));

                // Perform BPE merges
                for (int i = 0; i < num_merges; ++i) {
                    auto stats = get_stats(ids);
                    if (stats.empty()) break;

                    // Find the most frequent pair that does not exceed BPE_TOKENIZER_MAX_TOKEN_LENGTH
                    auto pair = get_most_frequent_pair(stats);

                    // Check if the resulting token would exceed BPE_TOKENIZER_MAX_TOKEN_LENGTH
                    size_t new_token_length = vocab[pair.first].size() + vocab[pair.second].size();
                    if (new_token_length > BPE_TOKENIZER_MAX_TOKEN_LENGTH) {
                        if (verbose)
                        {
                            std::cout << "\r"
                                << std::setw(100) << std::flush
                                << "\rskipping merge " << std::to_string(i + 1) << "/" << std::to_string(num_merges) << ": ("
                                << std::to_string(pair.first) << "," << std::to_string(pair.second) << ") -> new token length "
                                << std::to_string(new_token_length) << " exceeds limit of " << std::to_string(BPE_TOKENIZER_MAX_TOKEN_LENGTH)
                                << std::flush;
                        }
                        continue; // Skip this merge
                    }

                    int idx = (BPE_TOKENIZER_BASE_VOCAB_SIZE + (int)special_tokens.size()) + i;
                    ids = merge(ids, pair, idx);
                    merges[pair] = idx;
                    vocab[idx].insert(vocab[idx].end(), vocab[pair.first].begin(), vocab[pair.first].end());
                    vocab[idx].insert(vocab[idx].end(), vocab[pair.second].begin(), vocab[pair.second].end());

                    if (verbose)
                    {
                        std::cout << "\r"
                            << std::setw(100) << std::flush
                            << "\rmerge " << std::to_string(i + 1) << "/" << std::to_string(num_merges) << ": ("
                            << std::to_string(pair.first) << "," << std::to_string(pair.second) << ") -> " << std::to_string(idx)
                            << " (" << bytes_to_string(vocab[idx]) << ") had "
                            << std::to_string(stats[pair]) << " occurrences"
                            << std::endl;
                    }
                }
            }

            // Encode the given text into subword tokens
            std::vector<int> encode(const std::string& text) const
            {
                std::vector<int> result_ids;
                std::mutex result_mutex;

                // Split the text into paragraphs based on newline characters
                std::vector<std::string> paragraphs;
                size_t start = 0, end = text.find('\n');
                while (end != std::string::npos) {
                    std::string paragraph = text.substr(start, end - start);
                    if (!paragraph.empty()) paragraphs.push_back(paragraph);
                    start = end + 1;
                    end = text.find('\n', start);
                }
                // Add the last paragraph (if any) and only if it's not empty
                if (start < text.size()) {
                    std::string paragraph = text.substr(start);
                    if (!paragraph.empty()) paragraphs.push_back(paragraph);
                }

                // Function to encode a single paragraph
                auto encode_paragraph = [this](const std::string& paragraph) -> std::vector<int> {
                    std::vector<int> ids;
                    ids.reserve(paragraph.size());
                    for (char c : paragraph) ids.push_back(static_cast<uint8_t>(c));

                    auto stats = get_stats(ids);
                    std::priority_queue<std::pair<int, std::pair<int, int>>> pq;
                    for (const auto& stat : stats) {
                        const std::pair<int, int>& pair = stat.first;
                        if (merges.count(pair)) pq.push({ merges.at(pair), pair });
                    }

                    while (!pq.empty()) {
                        const auto& top_element = pq.top();
                        const std::pair<int, int>& pair = top_element.second;
                        pq.pop();

                        bool pair_found = false;
                        for (size_t i = 0; i < ids.size() - 1; ++i) {
                            if (ids[i] == pair.first && ids[i + 1] == pair.second) {
                                pair_found = true;
                                break;
                            }
                        }
                        if (!pair_found) continue;

                        int idx = merges.at(pair);
                        ids = merge(ids, pair, idx);

                        stats = get_stats(ids);
                        for (const auto& stat : stats) {
                            const std::pair<int, int>& new_pair = stat.first;
                            if (merges.count(new_pair)) pq.push({ merges.at(new_pair), new_pair });
                        }
                    }

                    return ids;
                };

                // Special case: if there's only one paragraph, no need for threads
                int sot_tok = get_special_token_id("<text>");
                int eot_tok = get_special_token_id("</text>");
                if (paragraphs.size() == 1) {
                    std::vector<int> paragraph_ids = encode_paragraph(paragraphs[0]);
                    result_ids.push_back(sot_tok);
                    result_ids.insert(result_ids.end(), paragraph_ids.begin(), paragraph_ids.end());
                    result_ids.push_back(eot_tok);
                    return result_ids;
                }

                // Launch encoding tasks in parallel for multiple paragraphs
                std::vector<std::future<std::vector<int>>> futures;
                for (const auto& paragraph : paragraphs)
                    futures.push_back(std::async(std::launch::async, encode_paragraph, paragraph));

                // Collect results in order
                for (auto& future : futures) {
                    std::vector<int> paragraph_ids = future.get();
                    std::lock_guard<std::mutex> lock(result_mutex);
                    result_ids.push_back(sot_tok);
                    result_ids.insert(result_ids.end(), paragraph_ids.begin(), paragraph_ids.end());
                    result_ids.push_back(eot_tok);
                }
                return result_ids;
            }

            // Decode a single token ID back into text
            std::string decode(int id, bool display_special_tokens = true) const
            {
                return decode(std::vector<int>({ id }), display_special_tokens);
            }

            // Decode a sequence of token IDs back into text
            std::string decode(const std::vector<int>& ids, bool display_special_tokens = true) const
            {
                std::vector<uint8_t> bytes;
                int vocab_size = static_cast<int>(get_vocab_size());
                for (int id : ids)
                {
                    if (id < vocab_size)
                    {
                        // Check if the ID is a special token
                        auto it = special_token_map.find(id);
                        if (it != special_token_map.end())
                        {
                            // It's a special token, get the corresponding string
                            if (display_special_tokens) bytes.insert(bytes.end(), it->second.begin(), it->second.end());
                        }
                        else
                        {
                            // It's a regular token, get the bytes from the vocabulary
                            auto& token = vocab.at(id);
                            bytes.insert(bytes.end(), token.begin(), token.end());
                        }
                    }
                }
                return std::string(bytes.begin(), bytes.end());
            }

            // Save the tokenizer model and vocabulary to file
            friend void serialize(const bpe_tokenizer& tok, std::ostream& out)
            {
                serialize("bpe_tokenizer2_", out);
                serialize(tok.special_tokens, out);
                serialize(tok.special_token_map, out);
                serialize(tok.merges, out);
                serialize(tok.vocab, out);
                serialize(tok.vocab_size, out);
            }

            // Load the tokenizer model and vocabulary from file
            friend void deserialize(bpe_tokenizer& tok, std::istream& in) {
                std::string version;
                dlib::deserialize(version, in);
                if (version != "bpe_tokenizer2_")
                    throw dlib::serialization_error("Unexpected version '" + version + "' found while deserializing dlib::bpe_tokenizer_.");
                deserialize(tok.special_tokens, in);
                deserialize(tok.special_token_map, in);
                deserialize(tok.merges, in);
                deserialize(tok.vocab, in);
                deserialize(tok.vocab_size, in);
            }

            // Get the ID of a special token
            int get_special_token_id(const std::string& token) const
            {
                auto it = special_tokens.find(token);
                if (it != special_tokens.end()) return it->second;
                throw std::runtime_error("Special token not found: " + token);
            }

            // Get the total vocabulary size
            size_t get_vocab_size() const
            {
                return (vocab.size() + special_tokens.size());
            }

        private:
            std::map<std::string, int> special_tokens;
            std::unordered_map<int, std::string> special_token_map;
            std::map<std::pair<int, int>, int> merges;
            std::map<int, std::vector<uint8_t>> vocab;
            int vocab_size;

            // Get frequency statistics of adjacent token pairs
            struct pair_hash {
                template <class T1, class T2>
                std::size_t operator()(const std::pair<T1, T2>& p) const
                {
                    auto hash1 = std::hash<T1>{}(p.first);
                    auto hash2 = std::hash<T2>{}(p.second);
                    return hash1 ^ (hash2 << 1);
                }
            };
            std::unordered_map<std::pair<int, int>, int, pair_hash> get_stats(const std::vector<int>& ids) const
            {
                std::unordered_map<std::pair<int, int>, int, pair_hash> global_stats;
                std::mutex global_stats_mutex;

                auto worker = [&](size_t start, size_t end) {
                    std::unordered_map<std::pair<int, int>, int, pair_hash> local_stats;
                    for (size_t i = start; i < end - 1 && i + 1 < ids.size(); ++i)
                        local_stats[{ids[i], ids[i + 1]}]++;

                    std::lock_guard<std::mutex> lock(global_stats_mutex);
                    for (const auto& pair : local_stats)
                        global_stats[pair.first] += pair.second;
                };

                size_t num_threads = std::thread::hardware_concurrency();
                size_t segment_size = ids.size() / num_threads;
                std::vector<std::thread> threads;

                for (size_t t = 0; t < num_threads; ++t)
                {
                    size_t start = t * segment_size;
                    size_t end = (t == num_threads - 1) ? ids.size() : start + segment_size;
                    threads.emplace_back(worker, start, end);
                }

                for (auto& thread : threads) thread.join();

                return global_stats;
            }

            // Finds the most frequent pair of tokens in the given statistics map that does not exceed the maximum token length
            std::pair<int, int> get_most_frequent_pair(const std::unordered_map<std::pair<int, int>, int, pair_hash>& stats) const
            {
                std::pair<int, int> best_pair = { -1, -1 }; // Initialize the best pair to an invalid value
                double max_score = 0; // Initialize the maximum score to 0

                // Iterate over all pairs in the statistics map
                for (const auto& stat : stats) {
                    const std::pair<int, int>& pair = stat.first; // Extract the token pair
                    int count = stat.second; // Extract the frequency count

                    // Check if the new token formed by merging the pair would exceed the maximum allowed length
                    size_t new_token_length = vocab.at(pair.first).size() + vocab.at(pair.second).size();
                    if (new_token_length > BPE_TOKENIZER_MAX_TOKEN_LENGTH) continue; // Skip this pair if it exceeds the maximum token length

                    // Calculate the score for this pair (frequency * length_penalty)
                    double score = (size_t)count * (new_token_length > (BPE_TOKENIZER_MAX_TOKEN_LENGTH / 2) ? 1.75 : 1.0);

                    // Update the best pair if the current pair has a higher score
                    if (score > max_score)
                    {
                        best_pair = pair;
                        max_score = score;
                    }
                }

                return best_pair; // Return the pair with the highest score
            }

            // Merge the most frequent pair in the token sequence
            std::vector<int> merge(std::vector<int>& ids, const std::pair<int, int>& pair, int idx) const
            {
                std::vector<int> new_ids;
                new_ids.reserve(ids.size()); // Reserve space to avoid reallocations

                for (size_t i = 0; i < ids.size(); ++i)
                {
                    if (i < ids.size() - 1 && ids[i] == pair.first && ids[i + 1] == pair.second)
                    {
                        new_ids.push_back(idx); // Replace the pair with the new token ID
                        i++; // Skip the next token
                    }
                    else new_ids.push_back(ids[i]); // Keep the current token
                }

                return new_ids;
            }

            static std::string base64_encode(const std::string& input) {
                dlib::base64 encoder;
                std::istringstream sin(input);
                std::ostringstream sout;
                encoder.encode(sin, sout);
                return sout.str();
            }

            // Convert a sequence of bytes to a readable string
            static std::string bytes_to_string(const std::vector<uint8_t>& bytes)
            {
                std::string data(bytes.begin(), bytes.end());
                return base64_encode(data);
            }

        };

    }

    #endif // DLIB_BPE_TOKENIZER_H