#include 
#include 
#include 
#include 

#include 
#include 

std::vector random_embeddings(int rows, int dimensions) {
    std::random_device rd;
    std::mt19937_64 prng{rd()};
    std::uniform_real_distribution dist{0, 1};

    std::vector embeddings;
    embeddings.reserve(rows);
    for (int i = 0; i < rows; i++) {
        std::vector embedding;
        embedding.reserve(dimensions);
        for (int j = 0; j < dimensions; j++) {
            embedding.push_back(dist(prng));
        }
        embeddings.push_back(embedding);
    }
    return embeddings;
}

std::vector random_categories(int rows) {
    std::random_device rd;
    std::mt19937_64 prng{rd()};
    std::uniform_int_distribution dist{1, 100};

    std::vector categories;
    categories.reserve(rows);
    for (int i = 0; i < rows; i++) {
        categories.push_back(dist(prng));
    }
    return categories;
}

int main() {
    // generate random data
    int rows = 100000;
    int dimensions = 128;
    std::vector embeddings = random_embeddings(rows, dimensions);
    std::vector categories = random_categories(rows);
    std::vector queries = random_embeddings(10, dimensions);

    // enable extensions
    pqxx::connection conn{"dbname=pgvector_citus"};
    pqxx::nontransaction tx{conn};
    tx.exec("CREATE EXTENSION IF NOT EXISTS citus");
    tx.exec("CREATE EXTENSION IF NOT EXISTS vector");

    // GUC variables set on the session do not propagate to Citus workers
    // https://github.com/citusdata/citus/issues/462
    // you can either:
    // 1. set them on the system, user, or database and reconnect
    // 2. set them for a transaction with SET LOCAL
    tx.exec("ALTER DATABASE pgvector_citus SET maintenance_work_mem = '512MB'");
    tx.exec("ALTER DATABASE pgvector_citus SET hnsw.ef_search = 20");
    conn.close();

    // reconnect for updated GUC variables to take effect
    pqxx::connection conn2{"dbname=pgvector_citus"};
    pqxx::nontransaction tx2{conn2};

    std::cout