#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <fstream>
#include <iostream>
#include <mutex>
#include <optional>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#ifdef _OPENMP
#include <omp.h>
#else
inline int omp_get_max_threads(){ return 1; }
#endif

using namespace std;

static vector<string> tokenize_whitespace(const string &s){
    istringstream iss(s);
    vector<string> out;
    string t;
    while (iss >> t) out.push_back(t);
    return out;
}

using StrPtr = const string*;
using TokenId = uint32_t;
static constexpr TokenId TOKEN_ID_INVALID = 0xFFFFFFFFu;

static inline size_t popcount_u64(uint64_t x){
#if defined(_MSC_VER)
    return static_cast<size_t>(__popcnt64(x));
#else
    return static_cast<size_t>(__builtin_popcountll(static_cast<unsigned long long>(x)));
#endif
}

struct StringInterner {
    unordered_set<string> pool;
    unordered_map<string, TokenId> id_by_string;
    vector<const string*> string_by_id;
    mutable mutex m;

    const string* intern(const string &s){
        lock_guard<mutex> lk(m);

        auto [it, inserted] = pool.emplace(s);
        if (inserted){
            const TokenId id = static_cast<TokenId>(string_by_id.size());
            id_by_string.emplace(*it, id);
            string_by_id.push_back(&*it);
        }

        return &*it;
    }

    TokenId id_of(const string &s) const {
        lock_guard<mutex> lk(m);

        auto it = id_by_string.find(s);
        return (it == id_by_string.end()) ? TOKEN_ID_INVALID : it->second;
    }

    TokenId id_of(const string *p) const {
        return p ? id_of(*p) : TOKEN_ID_INVALID;
    }

    size_t size() const {
        lock_guard<mutex> lk(m);
        return string_by_id.size();
    }
};

static inline void bitset_set(uint64_t *bits, size_t words, TokenId id){
    if (id == TOKEN_ID_INVALID) return;
    const size_t idx = static_cast<size_t>(id >> 6);
    if (idx >= words) return;
    bits[idx] |= (1ULL << (id & 63u));
}

static inline size_t bitset_count(const uint64_t *bits, size_t words){
    size_t total = 0;
    for (size_t i = 0; i < words; ++i) total += popcount_u64(bits[i]);
    return total;
}

using NextSet = vector<StrPtr>;

struct NgramHash {
    size_t operator()(const vector<StrPtr>& v) const noexcept {
        size_t seed = v.size();
        for(auto& p : v) {
            seed ^= hash<StrPtr>()(p) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
        }
        return seed;
    }
};

struct KnowledgeBase {
    StringInterner interner;
    // Modified to use a vector of StrPtr (N-gram) as the key
    unordered_map<vector<StrPtr>, NextSet, NgramHash> next;
    mutable mutex m;

    void add_ngram(const vector<StrPtr>& ctx, StrPtr v){
        lock_guard<mutex> lk(m);
        auto &vec = next[ctx];
        for (auto p : vec) if (p == v) return;
        vec.push_back(v);
    }

    optional<NextSet> lookup_ngram(const vector<StrPtr>& ctx) const {
        lock_guard<mutex> lk(m);
        auto it = next.find(ctx);
        if (it == next.end()) return nullopt;
        return it->second;
    }

    // Fast read-only lookup returning direct pointer to avoid mutex locking and vector copies
    const NextSet* lookup_ngram_fast(const vector<StrPtr>& ctx) const {
        auto it = next.find(ctx);
        return (it != next.end()) ? &(it->second) : nullptr;
    }
};

static void learn_tokens_ngram(KnowledgeBase &kb, const vector<string>& tokens, int n_gram_size) {
    vector<StrPtr> window;
    for (const auto& tok : tokens) {
        StrPtr tp = kb.interner.intern(tok);
        for (size_t j = 0; j < window.size(); ++j) {
            vector<StrPtr> ctx(window.begin() + j, window.end());
            kb.add_ngram(ctx, tp);
        }
        window.push_back(tp);
        if (window.size() > static_cast<size_t>(n_gram_size)) window.erase(window.begin());
    }
}

static vector<StrPtr>
intern_tokens(KnowledgeBase &kb, const vector<string> &tokens)
{
    vector<StrPtr> out;
    out.reserve(tokens.size());
    for (const auto &t : tokens) out.push_back(kb.interner.intern(t));
    return out;
}


static const NextSet* lookup_ngram_with_backoff(
    const KnowledgeBase &kb,
    const vector<StrPtr> &ctx_path,
    int n_gram_size)
{
    int cur_len = min(static_cast<int>(ctx_path.size()), n_gram_size);
    while (cur_len > 0) {
        vector<StrPtr> window(ctx_path.end() - cur_len, ctx_path.end());
        const NextSet* res = kb.lookup_ngram_fast(window);
        if (res && !res->empty()) {
            return res;
        }
        cur_len--;
    }
    return nullptr;
}

struct CandState {
    StrPtr orig_cand;
    vector<vector<StrPtr>> frontier;
    double total_score;
};

static string choose_candidate(
    const NextSet &cands,
    const vector<StrPtr> &history_ptrs,
    const KnowledgeBase &kb,
    const unordered_map<string, int> &recent_counts,
    double repeat_penalty,
    int n_gram_size)
{
    if (cands.empty()) return string();
    if (cands.size() == 1) return *cands[0];

    // 1. Build historical context bitset once
    const size_t words = max<size_t>(1, (kb.interner.size() + 63u) / 64u);
    vector<uint64_t> agg_words(words, 0ULL);

    for (StrPtr t : history_ptrs) {
        TokenId tid = kb.interner.id_of(t);
        if (tid != TOKEN_ID_INVALID) bitset_set(agg_words.data(), words, tid);
    }
    const size_t agg_count = bitset_count(agg_words.data(), words);

    // Build base context tail from history (up to N-1 tokens)
    vector<StrPtr> base_ctx;
    int base_len = min(static_cast<int>(history_ptrs.size()), n_gram_size - 1);
    if (base_len > 0) {
        base_ctx.assign(history_ptrs.end() - base_len, history_ptrs.end());
    }

    // 2. Initialize iterative state for ALL original candidates
    vector<CandState> tied_states(cands.size());
    for (size_t i = 0; i < cands.size(); ++i) {
        tied_states[i].orig_cand = cands[i];

        // Initial frontier path for candidate = base_ctx + candidate
        vector<StrPtr> init_path = base_ctx;
        init_path.push_back(cands[i]);
        if (init_path.size() > static_cast<size_t>(n_gram_size)) {
            init_path.erase(init_path.begin(), init_path.end() - n_gram_size);
        }

        tied_states[i].frontier.push_back(move(init_path));
        tied_states[i].total_score = 0.0;
    }

    // 3. Iterative deepening BFS up to n_gram_size across all candidates
    for (int depth = 1; depth <= n_gram_size; ++depth) {

        #ifdef _OPENMP
        #pragma omp parallel for schedule(dynamic)
        #endif
        for (ptrdiff_t i = 0; i < static_cast<ptrdiff_t>(tied_states.size()); ++i) {
            double depth_max_score = -1.0;
            vector<vector<StrPtr>> next_frontier;

            for (const auto &path : tied_states[i].frontier) {
                // Perform n-gram backoff lookup using the full path context
                const NextSet* opt_next = lookup_ngram_with_backoff(kb, path, n_gram_size);
                if (!opt_next || opt_next->empty()) continue;

                size_t count = 0;
                size_t inter = 0;

                // O(1) bitwise lookups bypass allocating new bitsets
                for (StrPtr nxt : *opt_next) {
                    TokenId nid = kb.interner.id_of(nxt);
                    if (nid != TOKEN_ID_INVALID) {
                        count++;
                        if ((agg_words[nid >> 6] & (1ULL << (nid & 63u))) != 0) {
                            inter++;
                        }
                    }
                }

                const size_t uni = agg_count + count - inter;
                const double score = uni ? static_cast<double>(inter) / static_cast<double>(uni) : 0.0;

                if (score > depth_max_score) {
                    depth_max_score = score;
                }

                // Append new context paths to next frontier if under max depth
                if (depth < n_gram_size) {
                    for (StrPtr nxt : *opt_next) {
                        vector<StrPtr> next_path = path;
                        next_path.push_back(nxt);
                        if (next_path.size() > static_cast<size_t>(n_gram_size)) {
                            next_path.erase(next_path.begin(), next_path.end() - n_gram_size);
                        }
                        next_frontier.push_back(move(next_path));
                    }
                }
            }

            // Apply repetition penalty on the base token level at depth 1
            if (depth == 1) {
                const string &tok = *(tied_states[i].orig_cand);
                auto rc_it = recent_counts.find(tok);
                const int cnt = (rc_it == recent_counts.end()) ? 0 : rc_it->second;
                if (depth_max_score > -0.5) {
                    depth_max_score -= repeat_penalty * static_cast<double>(cnt);
                } else {
                    depth_max_score = -repeat_penalty * static_cast<double>(cnt);
                }
            } else {
                if (depth_max_score < 0.0) {
                    depth_max_score = 0.0;
                }
            }

            // Accumulate depth score into candidate's total score
            tied_states[i].total_score += depth_max_score;

            // Deduplicate next frontier to prevent exponential path expansion
            if (next_frontier.size() > 1) {
                sort(next_frontier.begin(), next_frontier.end());
                next_frontier.erase(unique(next_frontier.begin(), next_frontier.end()), next_frontier.end());
            }
            tied_states[i].frontier = move(next_frontier);
        }
    }

    // 4. Final Selection: Identify highest total score across all depths
    double best_total = -1e9;
    for (const auto &st : tied_states) {
        if (st.total_score > best_total) {
            best_total = st.total_score;
        }
    }

    // Collect all candidates tied for the best total score
    vector<StrPtr> tied_cands;
    for (const auto &st : tied_states) {
        if (st.total_score >= best_total - 1e-9) {
            tied_cands.push_back(st.orig_cand);
        }
    }

    // Break remaining ties alphabetically
    string best_tok = *(tied_cands[0]);
    for (size_t i = 1; i < tied_cands.size(); ++i) {
        if (*(tied_cands[i]) < best_tok) {
            best_tok = *(tied_cands[i]);
        }
    }

    return best_tok;
}

static vector<string> construct_response(KnowledgeBase &kb,
                                                   const vector<string> &prompt_toks,
                                                   size_t response_maxlen,
                                                   double repeat_penalty,
                                                   int n_gram_size,
                                                   bool allow_path_loops)
{
    vector<string> resp;
    if (prompt_toks.empty() || response_maxlen == 0) return resp;

    auto prompt_ptrs = intern_tokens(kb, prompt_toks);
    unordered_map<string, int> recent_counts;

    // The context combines prompt and generated text
    vector<StrPtr> context = prompt_ptrs;

    auto would_create_2_cycle = [&](const string &cand) -> bool {
        if (resp.size() < 3) return false;
        return cand == resp[resp.size() - 2] &&
               resp.back() == resp[resp.size() - 3];
    };

    for (size_t step = 0; step < response_maxlen; ++step){
        NextSet candidates;
        bool found = false;

        int current_n = min(static_cast<int>(context.size()), n_gram_size);
        while (current_n > 0) {
            vector<StrPtr> search_ctx(context.end() - current_n, context.end());
            auto opt = kb.lookup_ngram(search_ctx);
            if (opt && !opt->empty()) {
                candidates = *opt;
                found = true;
                break;
            }
            current_n--;
        }

        if (!found || candidates.empty()) break;

        if (candidates.size() == 1){
            string only = *candidates[0];
            if (!allow_path_loops && recent_counts[only] > 0) break;
            resp.push_back(only);
            StrPtr ptr = kb.interner.intern(only);
            context.push_back(ptr);
            recent_counts[only] += 1;
            cout << only << ' ' << flush;
            continue;
        }

        string chosen = choose_candidate(
            candidates, context, kb, recent_counts, repeat_penalty, n_gram_size
        );

        if (chosen.empty()) break;
        if (!allow_path_loops && would_create_2_cycle(chosen)) break;

        resp.push_back(chosen);
        StrPtr chosen_ptr = kb.interner.intern(chosen);
        context.push_back(chosen_ptr);
        recent_counts[chosen] += 1;
        cout << chosen << ' ' << flush;
    }

    return resp;
}

static void learn_from_file(KnowledgeBase &kb, const string &fname, int n_gram_size){
    ifstream ifs(fname);
    if (!ifs) return;
    vector<string> tokens;
    string tok;
    while (ifs >> tok) tokens.push_back(tok);
    learn_tokens_ngram(kb, tokens, n_gram_size);
}

static void learn_files_parallel(KnowledgeBase &kb, const vector<string> &files, int n_gram_size){
    #pragma omp parallel for schedule(dynamic)
    for (ptrdiff_t i=0;i<static_cast<ptrdiff_t>(files.size());++i){
        learn_from_file(kb, files[(size_t)i], n_gram_size);
    }
}

static constexpr uint64_t KB_MAGIC   = 0x434850434B535641ULL;
static constexpr uint64_t KB_VERSION = 2ULL;

static void write_u64(ostream &os, uint64_t v) {
    do {
        unsigned char byte = v & 0x7F;
        v >>= 7;
        if (v != 0) byte |= 0x80;
        os.put(static_cast<char>(byte));
    } while (v != 0);
    if (!os) throw runtime_error("write_u64 failed");
}

static uint64_t read_u64(istream &is) {
    uint64_t v = 0;
    unsigned int shift = 0;
    unsigned char byte;
    do {
        if (!is.read(reinterpret_cast<char*>(&byte), 1)) {
            throw runtime_error("read_u64 failed");
        }
        v |= (static_cast<uint64_t>(byte & 0x7F) << shift);
        shift += 7;
    } while (byte & 0x80);
    return v;
}

static void write_string(ostream &os, const string &s){
    write_u64(os, static_cast<uint64_t>(s.size()));
    if (!s.empty()){
        os.write(s.data(), static_cast<streamsize>(s.size()));
        if(!os) throw runtime_error("write_string failed");
    }
}

static string read_string(istream &is){
    uint64_t n = read_u64(is);
    //if (n > (1ULL << 30)) throw runtime_error("save file is corrupted: string too large");

    string s;
    s.resize(static_cast<size_t>(n));

    if (n != 0){
        is.read(&s[0], static_cast<streamsize>(n));
        if(!is) throw runtime_error("read_string failed");
    }
    return s;
}

static void save_kb_binary(const KnowledgeBase &kb, const string &fname){
    const string temp = fname + ".tmp";

    {
        ofstream ofs(temp.c_str(), ios::binary | ios::trunc);
        if (!ofs) throw runtime_error("cannot open temp save file");

        vector<string> pool;
        pool.reserve(kb.interner.pool.size());
        for (const auto &s : kb.interner.pool) pool.push_back(s);

        sort(pool.begin(), pool.end());

        unordered_map<string, uint64_t> id;
        id.reserve(pool.size());
        for (uint64_t i = 0; i < static_cast<uint64_t>(pool.size()); ++i){
            id.emplace(pool[(size_t)i], i);
        }

        write_u64(ofs, KB_MAGIC);
        write_u64(ofs, KB_VERSION);

        write_u64(ofs, static_cast<uint64_t>(pool.size()));
        for (const auto &s : pool) write_string(ofs, s);

        write_u64(ofs, static_cast<uint64_t>(kb.next.size()));
        for (const auto &pr : kb.next){
            write_u64(ofs, static_cast<uint64_t>(pr.first.size()));
            for (StrPtr ctx_tok : pr.first){
                write_u64(ofs, id.at(*ctx_tok));
            }
            write_u64(ofs, static_cast<uint64_t>(pr.second.size()));
            for (StrPtr nxt : pr.second){
                write_u64(ofs, id.at(*nxt));
            }
        }

        ofs.flush();
        if (!ofs) throw runtime_error("failed while writing temp save file");
    }

    remove(fname.c_str());
    if (rename(temp.c_str(), fname.c_str()) != 0){
        remove(temp.c_str());
        throw runtime_error("failed to commit save file");
    }
}

static void load_kb_binary(KnowledgeBase &kb, const string &fname){
    ifstream ifs(fname, ios::binary);
    if (!ifs) throw runtime_error("cannot open load file");

    const uint64_t magic = read_u64(ifs);
    if (magic != KB_MAGIC) throw runtime_error("bad save file magic");

    const uint64_t version = read_u64(ifs);
    if (version != KB_VERSION) throw runtime_error("unsupported save file version");

    const uint64_t N = read_u64(ifs);
    //if (N > (1ULL << 26)) throw runtime_error("save file is corrupted: pool too large");

    vector<string> strings;
    strings.reserve(static_cast<size_t>(N));

    for (uint64_t i = 0; i < N; ++i){
        strings.push_back(read_string(ifs));
    }

    kb.interner.pool.clear();
    kb.interner.pool.reserve(static_cast<size_t>(N));

    vector<StrPtr> ptrs;
    ptrs.reserve(static_cast<size_t>(N));
    for (const auto &s : strings){
        ptrs.push_back(kb.interner.intern(s));
    }

    // Rebuild next
    const uint64_t E = read_u64(ifs);
    //if (E > (1ULL << 26)) throw runtime_error("save file is corrupted: graph too large");

    {
        lock_guard<mutex> lk(kb.m);
        kb.next.clear();
        kb.next.reserve(static_cast<size_t>(E));
    }

    for (uint64_t i = 0; i < E; ++i){
        const uint64_t ctx_size = read_u64(ifs);
        //if (ctx_size > (1ULL << 16)) throw runtime_error("save file is corrupted: n-gram context too large");

        vector<StrPtr> ctx;
        ctx.reserve(static_cast<size_t>(ctx_size));
        for(uint64_t c = 0; c < ctx_size; ++c) {
            const uint64_t c_idx = read_u64(ifs);
            if (c_idx >= ptrs.size()) throw runtime_error("save file is corrupted: bad context key");
            ctx.push_back(ptrs[(size_t)c_idx]);
        }

        const uint64_t M = read_u64(ifs);
        //if (M > (1ULL << 26)) throw runtime_error("save file is corrupted: graph degree too large");

        NextSet vec;
        vec.reserve(static_cast<size_t>(M));

        for (uint64_t j = 0; j < M; ++j){
            const uint64_t v_idx = read_u64(ifs);
            if (v_idx >= ptrs.size()) throw runtime_error("save file is corrupted: bad graph value");
            vec.push_back(ptrs[(size_t)v_idx]);
        }

        {
            lock_guard<mutex> lk(kb.m);
            kb.next.emplace(move(ctx), move(vec));
        }
    }
}

static void print_kb_info(const string &fname) {
    ifstream ifs(fname, ios::binary);
    if (!ifs) {
        cerr << "Error: cannot open knowledge-base file " << fname << "\n";
        return;
    }

    cout << "Reckoning knowledge-base " << fname << " info. ";

    try {
        const uint64_t magic = read_u64(ifs);
        if (magic != KB_MAGIC) throw runtime_error("bad save file magic");

        const uint64_t version = read_u64(ifs);
        if (version != KB_VERSION) throw runtime_error("unsupported save file version");

        const uint64_t N = read_u64(ifs);
        //if (N > (1ULL << 26)) throw runtime_error("save file is corrupted: pool too large");

        // Dynamically read and skip strings since lengths are now variable-length encoded
        for (uint64_t i = 0; i < N; ++i){
            uint64_t len = read_u64(ifs);
            //if (len > (1ULL << 30)) throw runtime_error("save file is corrupted: string too large");
            if (len > 0) ifs.seekg(static_cast<streamoff>(len), ios_base::cur);
        }

        const uint64_t E = read_u64(ifs);
        //if (E > (1ULL << 26)) throw runtime_error("save file is corrupted: graph too large");

        uint64_t max_ngram = 0;
        for (uint64_t i = 0; i < E; ++i){
            const uint64_t ctx_size = read_u64(ifs);
            //if (ctx_size > (1ULL << 16)) throw runtime_error("save file is corrupted: n-gram context too large");

            if (ctx_size > max_ngram) max_ngram = ctx_size;

            // Sequentially read (and skip) context keys using LEB128 stream parsing
            for(uint64_t c = 0; c < ctx_size; ++c) {
                read_u64(ifs);
            }

            const uint64_t M = read_u64(ifs);
            //if (M > (1ULL << 26)) throw runtime_error("save file is corrupted: graph degree too large");

            // Sequentially read (and skip) next-node values
            for(uint64_t j = 0; j < M; ++j){
                read_u64(ifs);
            }
        }

        cout << "The " << fname << " n-gram max size is " << max_ngram << "\n";

    } catch (const exception &e) {
        cerr << "Error reading knowledge-base info: " << e.what() << "\n";
    }
}

static void print_commands(const char *p){
    cout << p << " command-line interface options:\n";
    cout << "  --response-max-length N     Set maximum number of token(s) in a response. Default number is 500\n";
    cout << "  --save-kb FILE              Save the knowledge-base graph to a binary file.\n";
    cout << "  --load-kb FILE              Load a previously saved knowledge-base graph from a binary file.\n";
    cout << "  --info-kb FILE              Show the dictionary-depth and n-gram max size of a knowledge-base file.\n";
    cout << "  --static-kb                 Disable learning prompt(s) and response(s) during interactive session.\n";
    cout << "  --n-gram-max-size N         Set maximum size of the n-gram where N is the size. Default size is 3\n";
    cout << "  --learn f1 f2 ...           Learn from text file(s) to update the knowledge-base graph.\n";
    cout << "  --repeat-penalty P          Set penalty for repeated tokens in response. Default penalty is 0.7 (higher penalty reduce repetition).\n";
    cout << "  --allow-path-loops          Allow finite loops when the knowledge-base graph contains a circular single-path reference.\n";
    cout << "  --activate-agi              Activate the Artificial General Intelligence (AGI) features.\n";
    cout << "  --help                      Show " << p << " command-line interface options.\n";
}

int main(int argc, char **argv){
    size_t response_maxlen = 500;
    string savefile;
    string load_kb;
    int n_gram_size = 3;
    double repeat_penalty = 0.7; // default λ
    vector<string> learn_files;
    bool learn_chat = true;
    bool allow_path_loops = false;

    for (int i=1;i<argc;++i){
        string a = argv[i];
        if (a=="--help"){ print_commands(argv[0]); return 0; }
        if (a=="--response-max-length" && i+1<argc){ response_maxlen = stoul(argv[++i]); continue; }
        if (a=="--save-kb" && i+1<argc){ savefile = argv[++i]; continue; }
        if (a=="--load-kb" && i+1<argc){ load_kb = argv[++i]; continue; }
        if (a=="--info-kb" && i+1<argc){ print_kb_info(argv[++i]); return 0; }
        if (a=="--static-kb"){ learn_chat = false; continue; }
        if (a=="--n-gram-max-size" && i+1<argc){ n_gram_size = max(1, stoi(argv[++i])); continue; }
        if (a=="--learn"){ while(i+1<argc && argv[i+1][0] != '-') learn_files.push_back(argv[++i]); continue; }
        if (a=="--repeat-penalty" && i+1<argc){ repeat_penalty = stod(argv[++i]); continue; }
        if (a=="--allow-path-loops") { allow_path_loops = true; continue; }
        if (a=="--activate-agi"){
                cerr << "AGI subscription fee is $1,000 CAD per month. To pay the fee, contact Caleb Nwokocha via email cnwokocha@proton.me\n";
                return 0;
        }

        cerr << "Error: Unknown or incomplete command-line option '" << a << "'\n";
        print_commands(argv[0]);
        return 1;
    }

    KnowledgeBase kb;

    if (!load_kb.empty()){
        try { print_kb_info(load_kb);
        cerr << "Loading knowledge base " << load_kb << " using values from " << argv[0] << " command-line interface options: --n-gram-max-size " << n_gram_size <<
        ", --response-max-length " << response_maxlen << ", and --repeat-penalty " << repeat_penalty << "\n";
        cerr << "To change the value(s), use " << argv[0] << " command-line interface option(s).\n";
        load_kb_binary(kb, load_kb); cerr << "Loaded knowledge base: " << load_kb << "\n"; }
        catch (const exception &e){ cerr << "Error: " << e.what() << "\n"; }
    }

    if (!learn_files.empty()){
        cerr << "Learning from file/s (" << learn_files.size() << ") using threads=" << omp_get_max_threads() << "\n";
        learn_files_parallel(kb, learn_files, n_gram_size);
    }

    string line;
    cout << "Ready. Enter prompts.\n";
    while (cout << "> " , getline(cin, line)){
        if (line.empty()){ cout << "\n"; continue; }
        auto prompt_toks = tokenize_whitespace(line);
        if (learn_chat) {learn_tokens_ngram(kb, prompt_toks, n_gram_size);}
        auto resp = construct_response(kb, prompt_toks, response_maxlen, repeat_penalty, n_gram_size, allow_path_loops);
        cout << "\n";
        if (!resp.empty() && learn_chat){learn_tokens_ngram(kb, resp, n_gram_size);}
        if (!savefile.empty()){
            try { cerr << "Saving knowledge base: " << savefile << "\n";
            save_kb_binary(kb, savefile); cerr << "Saved knowledge base: " << savefile << "\n"; }
            catch (const exception &e){ cerr << "Error: " << e.what() << "\n"; }
        }
    }

    return 0;
}