#include <iostream>
#include <list>
#include <unordered_map>

class LFUCache
{
public:
    LFUCache(int capacity)
    {
        _capacity = capacity;
    }

    int get(int key)
    {
        if (!_keyToIterators.count(key))
            return -1;

        updateUsage(key);

        return *_keyToIterators[key].second;
    }

    void put(int key, int value)
    {
        if (_keyToIterators.count(key))
        {
            updateUsage(key);
            *_keyToIterators[key].second = value;
        }
        else
        {
            if (_frequencies.size() == _capacity)
            {
                _keyToIterators.erase(); // TODO: remove the key of LFU & LRU key,val.
                _frequencies.front().values.pop_back();
            }
            if (!_frequencies.empty() && _frequencies.front().values.empty())
                _frequencies.pop_front();
            if (!_frequencies.empty() && 1 == _frequencies.front().frequency)
                _frequencies.front().values.push_front(value);
            else
                _frequencies.push_front({ .frequency = 1, .values = {value} });
            _keyToIterators[key] = { _frequencies.begin(),
                                    _frequencies.begin()->values.begin() };
            // std::cout << "_frequencies.begin()->frequency = " <<  _frequencies.begin()->frequency << std::endl;
            // std::cout << "_frequencies.begin().front().values.begin() = " <<  *(_frequencies.begin() << std::endl;
        }
    }

private:
    void updateUsage(int key)
    {
        auto [elementFreqIt, elementIt] = _keyToIterators[key];
        // std::pair<std::list<Item>::iterator, std::list<int>::iterator> p = ;

        // std::list<Item>::iterator elementFreqIt = _keyToIterators[key].first;
        // std::list<int>::iterator elementIt = _keyToIterators[key].second;

        auto nextFreqIt = std::next(elementFreqIt);
        int value = *elementIt;
        int currentFreq = elementFreqIt->frequency;

        elementFreqIt->values.erase(elementIt);

        if (_frequencies.end() == nextFreqIt ||
            nextFreqIt->frequency - 1 != currentFreq)
        {
            auto insertIt = _frequencies.insert(
                nextFreqIt, { .frequency = currentFreq + 1, .values = {value} });
            _keyToIterators[key] = { insertIt, insertIt->values.begin() };
        }
        else
        {
            nextFreqIt->values.push_front(value);
            _keyToIterators[key] = { nextFreqIt, nextFreqIt->values.begin() };
        }

        if (elementFreqIt->values.empty())
            _frequencies.erase(elementFreqIt);
    }

public:
    struct Item
    {
        int frequency;
        std::list<int> values;
    };

    int _capacity;
    std::list<Item> _frequencies;
    std::unordered_map<
        int, std::pair<std::list<Item>::iterator, std::list<int>::iterator>>
        _keyToIterators;
};

int main(int argc, char** argv)
{
    LFUCache lfu(2);
    lfu.put(1, 1);                                           // cache=[1,_], cnt(1)=1
    lfu.put(2, 2);                                           // cache=[2,1], cnt(2)=1, cnt(1)=1
    std::cout << "lfu.get(1) = " << lfu.get(1) << std::endl; // return 1
                                                             // cache=[1,2], cnt(2)=1, cnt(1)=2
    lfu.put(3, 3);                                           // 2 is the LFU key because cnt(2)=1 is the smallest,
                                                             // invalidate 2. cache=[3,1], cnt(3)=1, cnt(1)=2
    std::cout << "lfu.get(2) = " << lfu.get(2) << std::endl; // return -1 (not found)
    std::cout << "lfu.get(3) = " << lfu.get(3) << std::endl; // return 3
                                                             // cache=[3,1], cnt(3)=2, cnt(1)=2
    lfu.put(4, 4);                                           // Both 1 and 3 have the same cnt, but 1 is LRU, invalidate 1.
                                                             // cache=[4,3], cnt(4)=1, cnt(3)=2
    std::cout << "lfu.get(1) = " << lfu.get(1) << std::endl; // return -1 (not found)
    std::cout << "lfu.get(3) = " << lfu.get(3) << std::endl; // return 3
                                                             // cache=[3,4], cnt(4)=1, cnt(3)=3
    std::cout << "lfu.get(4) = " << lfu.get(4) << std::endl; // return 4
                                                             // cache=[4,3], cnt(4)=2, cnt(3)=3
    return 0;
}