#include <random>

thread_local std::mt19937_64 rng(std::random_device{}());

void RandomXor(Int* currentKey, int bit_length, int flip_count) {
    alignas(32) uint64_t flipMasks[4] = {0};

    // Generate random positions and set bits in batches
    for (int i = 0; i < flip_count; ) {
        // Generate 4 random positions at once
        uint64_t rand_positions[4];
        for (int j = 0; j < 4 && i < flip_count; j++, i++) {
            rand_positions[j] = rng() % bit_length;
            int word = rand_positions[j] / 64;
            int bit = rand_positions[j] % 64;
            flipMasks[word] ^= (1ULL << bit);
        }
    }

    // Apply XOR in one AVX2 operation
    __m256i keyVec = _mm256_loadu_si256((__m256i*)currentKey->bits64);
    __m256i maskVec = _mm256_loadu_si256((__m256i*)flipMasks);
    __m256i result = _mm256_xor_si256(keyVec, maskVec);
    _mm256_storeu_si256((__m256i*)currentKey->bits64, result);

    // Clear masks to zero for next iteration
    memset(flipMasks, 0, sizeof(flipMasks));
}

void worker(Secp256K1* secp, int bit_length, int flip_count, int threadId, AVXCounter start, AVXCounter end) {
    const int fullBatchSize = 2 * POINTS_BATCH_SIZE;
    alignas(32) uint8_t localPubKeys[HASH_BATCH_SIZE][33];
    alignas(32) uint8_t localHashResults[HASH_BATCH_SIZE][20];
    alignas(32) int pointIndices[HASH_BATCH_SIZE];

    __m256i target16 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(TARGET_HASH160_RAW.data()));

    alignas(32) Point plusPoints[POINTS_BATCH_SIZE];
    alignas(32) Point minusPoints[POINTS_BATCH_SIZE];

    for (int i = 0; i < POINTS_BATCH_SIZE; i++) {
        Int tmp;
        tmp.SetInt32(i);
        plusPoints[i] = secp->ComputePublicKey(&tmp);
        minusPoints[i] = plusPoints[i];
        minusPoints[i].y.ModNeg();
    }

    alignas(32) Int deltaX[POINTS_BATCH_SIZE];
    IntGroup modGroup(POINTS_BATCH_SIZE);
    alignas(32) Int pointBatchX[fullBatchSize];
    alignas(32) Int pointBatchY[fullBatchSize];

    CombinationGenerator gen(bit_length, flip_count);
    gen.unrank(start.load());

    AVXCounter count;
    count.store(start.load());

    uint64_t actual_work_done = 0;
    auto last_report = chrono::high_resolution_clock::now();

    while (!stop_event.load() && count < end) {
    Int currentKey;
    currentKey.Set(&BASE_KEY);
    
    RandomXor(&currentKey, bit_length, flip_count);

    string keyStr = currentKey.GetBase16();
    keyStr = string(64 - keyStr.length(), '0') + keyStr;

     #pragma omp critical
    {
        g_threadPrivateKeys[threadId] = keyStr;
    }