#include <iostream>
#include <string>
#include <cstdlib>
#include <sys/mman.h>
#include <unistd.h>
#include <cstring>

int main(int argc, char** argv) {
    if (argc != 2) {
        std::cerr << "Usage: " << argv[0] << " <GB to allocate and lock>\n";
        return 1;
    }

    double gb = std::stod(argv[1]);
    size_t bytes = static_cast<size_t>(gb * 1024.0 * 1024.0 * 1024.0);

    std::cout << "Allocating " << gb << " GB (" << bytes << " bytes) of RAM...\n";
    
    void* ptr = mmap(NULL, bytes, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (ptr == MAP_FAILED) {
        perror("mmap failed");
        return 1;
    }

    std::cout << "Memory allocated. Faulting pages and pinning to RAM...\n";
    
    // Write to memory to map it to physical pages and prevent lazy allocation
    size_t page_size = sysconf(_SC_PAGESIZE);
    char* char_ptr = static_cast<char*>(ptr);
    for (size_t i = 0; i < bytes; i += page_size) {
        char_ptr[i] = 1;
    }

    // Mlock to pin it to RAM and prevent it from being swapped out itself
    if (mlock(ptr, bytes) != 0) {
        perror("mlock failed (you probably need to run with sudo)");
    } else {
        std::cout << "mlock successful.\n";
    }

    std::cout << "Memory is fully resident. Other inactive processes/caches should be pushed to swap.\n";
    std::cout << "Clearing filesystem caches...\n";
    
    int ret = system("echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null");
    if (ret != 0) {
        std::cerr << "Failed to clear caches. (Maybe sudo failed?)\n";
    } else {
        std::cout << "Caches cleared successfully.\n";
    }

    std::cout << "Unlocking and releasing memory...\n";
    munlock(ptr, bytes);
    munmap(ptr, bytes);

    std::cout << "Done! Try starting your model now.\n";
    return 0;
}