#include <iostream>
#include <chrono>
#include <vector>
#include <cassert>
#include <cstddef>

template <typename T>
class MemoryPool {
private:
    union Node {
        Node* next;
        char data[sizeof(T)];
    };

    Node* freeList;
    std::vector<Node*> blocks;
    size_t blockSize;
    size_t freeSlots;          // Added: free node count

    void allocateBlock() {
        Node* newBlock = static_cast<Node*>(::operator new(blockSize * sizeof(Node)));
        for (size_t i = 0; i < blockSize; ++i) {
            Node* node = newBlock + i;
            node->next = freeList;
            freeList = node;
        }
        freeSlots += blockSize;   // Add free nodes
        blocks.push_back(newBlock);
    }

public:
    explicit MemoryPool(size_t blockSize = 256) 
        : freeList(nullptr), blockSize(blockSize), freeSlots(0) {
        allocateBlock();
    }

    MemoryPool(const MemoryPool&) = delete;
    MemoryPool& operator=(const MemoryPool&) = delete;

    ~MemoryPool() {
        for (Node* block : blocks) {
            ::operator delete(block);
        }
    }

    T* allocate() {
        if (freeList == nullptr) {
            allocateBlock();
        }
        Node* node = freeList;
        freeList = node->next;
        --freeSlots;            // Decrement free count
        return reinterpret_cast<T*>(node);
    }

    void deallocate(T* ptr) {
        if (ptr == nullptr) return;
        Node* node = reinterpret_cast<Node*>(ptr);
        node->next = freeList;
        freeList = node;
        ++freeSlots;            // Increment free count
    }

    size_t freeCount() const { return freeSlots; }  // O(1)
};

struct TestObject {
    int id;
    double value;
    char padding[120];
};

void testWithMemoryPool(size_t numObjects) {
    MemoryPool<TestObject> pool(128);
    std::vector<TestObject*> objects;
    objects.reserve(numObjects);

    // 分配并构造
    for (size_t i = 0; i < numObjects; ++i) {
        TestObject* obj = pool.allocate();
        new (obj) TestObject{static_cast<int>(i), i * 0.5, {}};
        objects.push_back(obj);
        if ((i + 1) % (numObjects / 5) == 0)   // Progress indicator
            std::cout << "." << std::flush;
    }
    std::cout << " Allocation completed" << std::endl;

    // Validation
    for (size_t i = 0; i < numObjects; i += numObjects / 10) {
        assert(objects[i]->id == static_cast<int>(i));
        assert(objects[i]->value == i * 0.5);
    }

    // Destruct and deallocate
    for (TestObject* obj : objects) {
        obj->~TestObject();
        pool.deallocate(obj);
    }
}

void testWithNewDelete(size_t numObjects) {
    std::vector<TestObject*> objects;
    objects.reserve(numObjects);

    for (size_t i = 0; i < numObjects; ++i) {
        TestObject* obj = new TestObject{static_cast<int>(i), i * 0.5, {}};
        objects.push_back(obj);
        if ((i + 1) % (numObjects / 5) == 0)
            std::cout << "." << std::flush;
    }
    std::cout << " Allocation completed" << std::endl;

    for (size_t i = 0; i < numObjects; i += numObjects / 10) {
        assert(objects[i]->id == static_cast<int>(i));
        assert(objects[i]->value == i * 0.5);
    }

    for (TestObject* obj : objects) {
        delete obj;
    }
}

int main() {
    const size_t numObjects = 200000;  // 200,000 times, enough to show difference without lag
    const int repeatTimes = 5;

    std::cout << "=== Memory Pool Performance Test ===" << std::endl;
    std::cout << "Operations: " << numObjects << " allocations+deallocations (repeated " << repeatTimes << " times)" << std::endl;
    std::cout << "Object size: " << sizeof(TestObject) << " bytes" << std::endl;

    // Test memory pool
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < repeatTimes; ++i) {
        std::cout << "Memory pool round " << i+1 << " ";
        testWithMemoryPool(numObjects);
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto poolDuration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    std::cout << "Memory pool average time: " << poolDuration / repeatTimes << " ms" << std::endl;

    // Test new/delete
    start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < repeatTimes; ++i) {
        std::cout << "new/delete round " << i+1 << " ";
        testWithNewDelete(numObjects);
    }
    end = std::chrono::high_resolution_clock::now();
    auto newDuration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    std::cout << "new/delete average time: " << newDuration / repeatTimes << " ms" << std::endl;

    std::cout << "Speedup: " << static_cast<double>(newDuration) / poolDuration << "x" << std::endl;

    // Reuse test (no longer iterate the entire linked list, use O(1) freeCount instead)
    std::cout << "\n=== Memory Pool Reuse Test ===" << std::endl;
    MemoryPool<TestObject> pool(10);
    std::vector<TestObject*> ptrs;

    for (int i = 0; i < 13; ++i) {
        ptrs.push_back(pool.allocate());
        new (ptrs.back()) TestObject{i, i * 1.0, {}};
    }
    std::cout << "After allocating 13, free count: " << pool.freeCount() << std::endl;

    for (int i = 0; i < 5; ++i) {
        ptrs[i]->~TestObject();
        pool.deallocate(ptrs[i]);
    }
    std::cout << "After deallocating 5, free count: " << pool.freeCount() << std::endl;

    for (int i = 0; i < 5; ++i) {
        TestObject* obj = pool.allocate();
        new (obj) TestObject{100 + i, 99.9, {}};
        ptrs.push_back(obj);
    }
    std::cout << "After allocating 5 more, free count: " << pool.freeCount() << std::endl;

    for (TestObject* obj : ptrs) {
        obj->~TestObject();
        pool.deallocate(obj);
    }
    std::cout << "After deallocating all, free count: " << pool.freeCount() << std::endl;

    std::cout << "All tests passed!" << std::endl;
    return 0;
}