#include "SMS.h"
#include <random>

using namespace std;

void SMS::storeRequest(gp *payload)
{
    Thread thread = DramExtension::getExtension(payload).getThread();
    bool wasEmpty = isRequestBuffersEmpty();

    requestBuffers[thread].emplace_back(payload);

    if (inFlightMemRequestCounter.find(thread) == inFlightMemRequestCounter.end()) {
        inFlightMemRequestCounter[thread] = 0;
        cacheMisses[thread] = 0;
    }
    inFlightMemRequestCounter[thread]++;
    cacheMisses[thread]++;

    if (wasEmpty) {
        newRequest.notify(SC_ZERO_TIME);
    }
}

std::pair<Command, gp *> SMS::getNextRequest(Bank bank)
{
    if (bankBuffers[bank].empty()) {
        debugManager.printDebugMessage(name(),
                                       "Get next request on bank " + to_string(bank.ID()) + " : EMPTY buffer");
        return pair<Command, tlm::tlm_generic_payload *>(Command::NOP, NULL);
    } else {
        gp *payload = bankBuffers[bank].front();
        Command command = IScheduler::getNextCommand(*payload);
        if (command == Command::Read || command == Command::ReadA
                || command == Command::Write
                || command == Command::WriteA) {
            inFlightMemRequestCounter[DramExtension::getExtension(payload).getThread()]--;
            bankBuffers[bank].pop_front();
        }

        debugManager.printDebugMessage(name(),
                                       "Get next request on bank " + to_string(bank.ID()));
        return pair<Command, tlm::tlm_generic_payload *>(command, payload);
    }
}

void SMS::batchScheduler()
{
    sc_time memClk = Configuration::getInstance().memSpec.clk;
    std::default_random_engine generator;
    std::bernoulli_distribution distribution((double) SJFprobability / 100.0);

    while (true) {
        updateMPKCs(memClk);
        if (isRequestBuffersEmpty()) {
            wait(newRequest);
        } else {
            multiBatchFormation(memClk);
            if (existReadyBatches()) {
                if (!isSystemLightlyLoaded() && (existLowIntensityThread()
                                                 || distribution(generator))) {
                    pickSJF();
                } else {
                    pickRR();
                }
                drainOnePayloadFromReadybatch(memClk);
            } else {
                wait(memClk);
            }
        }
    }
}

/**
 * Pick a Thread according to Shortest-Job Policy
 * Save the picked one into lastSelectedThread
 * @return true if it can, otherwise false
 */
bool SMS::pickSJF()
{
    // find threads with ready batches
    std::vector<Thread> threadsWithReadyBatches;
    for (const auto &each : readyBatchInclusiveEndLocs) {
        if (!each.second.empty()) {
            // marked as thread with non-empty request buffer
            threadsWithReadyBatches.push_back(each.first);
        }
    }

    if (!threadsWithReadyBatches.empty()) {
        // pick shortest-job thread among threads with non-empty request buffer
        Thread &minThread = threadsWithReadyBatches.front();
        for (const auto &thread : threadsWithReadyBatches) {
            if (inFlightMemRequestCounter[thread] < inFlightMemRequestCounter[minThread]) {
                minThread = thread;
            }
        }

        // save selected thread
        lastSelectedThread = readyBatchInclusiveEndLocs.find(minThread);
        debugManager.printDebugMessage(name(),
                                       "[SJF] Select ready batch of thread " + to_string(minThread.ID()));
        return true;
    } else {
        return false;
    }
}

/**
 * Drain the picked request buffer into bank buffers
 * by move request one-by-one from start of the request buffer till last parameter
 * @param memClk
 * @param last
 */
void SMS::drainOnePayloadFromReadybatch(const sc_time &memClk)
{
    if (lastSelectedThread->second.empty()) {
        return;
    }
    const Thread &selectedThread = lastSelectedThread->first;

    const size_t &inclusiveEndLoc = lastSelectedThread->second.front();
    assert(inclusiveEndLoc < requestBuffers.size());

    for (size_t i = 0; i <= inclusiveEndLoc; ++i) {
        wait(memClk);
        Bank bank = DramExtension::getExtension(
                        requestBuffers[selectedThread].front()).getBank();
        bankBuffers[bank].emplace_back(requestBuffers[selectedThread].front());
        requestBuffers[selectedThread].pop_front();

        // decrement inclusive end locations of ready batches
        // except this ready batch
        for (size_t i = 1; i < readyBatchInclusiveEndLocs[selectedThread].size(); ++i) {
            --readyBatchInclusiveEndLocs[selectedThread][i];
        }

        debugManager.printDebugMessage(name(),
                                       "[SJF] Drain request in the ready batch of thread "
                                       + to_string((*lastSelectedThread).first.ID()) + " to bankbuffer "
                                       + to_string(bank.ID()));
    }
    lastSelectedThread->second.pop_front();
}

/**
 * Pick a Thread according to Round-Robin Policy
 * Save the picked one into lastSelectedThread
 * @return true if it can pick one, otherwise false
 */
bool SMS::pickRR()
{
    if (lastSelectedThread == readyBatchInclusiveEndLocs.end()) {
        lastSelectedThread = readyBatchInclusiveEndLocs.begin();
        if (!(*lastSelectedThread).second.empty()) {
            return true;
        }
    }

    std::map<Thread, std::deque<size_t>>::iterator savedOriginalNextSelectedThread =
                                          lastSelectedThread;

    do {
        lastSelectedThread++;
        if (lastSelectedThread == readyBatchInclusiveEndLocs.end()) {
            lastSelectedThread = readyBatchInclusiveEndLocs.begin();
        }
        if (lastSelectedThread == savedOriginalNextSelectedThread) {
            return false;
        }
    } while ((*lastSelectedThread).second.empty());

    debugManager.printDebugMessage(name(),
                                   "[RR] Select ready batch of thread " + to_string((
                                                                                        *lastSelectedThread).first.ID()));
    return true;
}

bool SMS::isSystemLightlyLoaded() const
{
    unsigned int totalRequest = 0;
    for (const auto &bankBuffer : bankBuffers) {
        totalRequest += bankBuffer.second.size();
    }
    return (totalRequest <= LOW_SYSTEM_LOAD);
}

bool SMS::existLowIntensityThread() const
{
    for (const auto &mpkcPerThread : MPKCs) {
        if (mpkcPerThread.second < LOW_MPKC) {
            return true;
        }
    }
    return false;
}

bool SMS::isThresholdAgeExceeded(const Thread &thread, sc_time const &memClk,
                                 size_t const &inclusiveBeginLoc, size_t const &exclusiveEndLoc)
{
    assert((exclusiveEndLoc - inclusiveBeginLoc) >= 1);
    // find the oldest request in the thread's batch
    sc_time oldestGenerationTime = sc_time_stamp();
    for (size_t i = inclusiveBeginLoc; i != exclusiveEndLoc; ++i) {
        sc_time reqGenerationTime = GenerationExtension::getExtension(
                                        requestBuffers[thread][i]).TimeOfGeneration();
        if (reqGenerationTime < oldestGenerationTime) {
            oldestGenerationTime = reqGenerationTime;
        }
    }

    // check threshold age according to the thread's MPKC
    sc_time oldestRequestAge = sc_time_stamp() - oldestGenerationTime;
    if ((MPKCs[thread] <= MEDIUM_MPKC)
            && (oldestRequestAge > (MEDIUM_THRESHOLD_AGE * memClk))) {
        return true;
    } else if ((MPKCs[thread] > MEDIUM_MPKC)
               && (oldestRequestAge > (HIGH_THRESHOLD_AGE * memClk))) {
        return true;
    } else {
        return false;
    }
}

void SMS::updateMPKCs(sc_time const &memClk)
{
    if (sc_time_stamp() % (MPKC_RESET_CYCLE * memClk) <= memClk) {
        // reset for every 10k clk cycles
        for (const auto &cacheMiss : cacheMisses) {
            MPKCs[cacheMiss.first] = 0;
        }
        debugManager.printDebugMessage(name(), "Reset MKKCs");
    } else {
        // update MPKC for every thread
        for (const auto &cacheMiss : cacheMisses) {
            MPKCs[cacheMiss.first] = (cacheMiss.second * 1000.0 * memClk) /
                                     (sc_time_stamp());
        }
        debugManager.printDebugMessage(name(), "Update MPKCs");
    }
}

bool SMS::isExceededReqBufferSize(size_t const &exclusiveEndLoc)
{
    return exclusiveEndLoc <= Configuration::getInstance().RequestBufferSize;
}

bool SMS::isRequestBuffersEmpty() const
{
    for (const auto &requestBuffer : requestBuffers) {
        if (!requestBuffer.second.empty()) {
            return false;
        }
    }
    return true;
}

bool SMS::existReadyBatches() const
{
    for (const auto &each : readyBatchInclusiveEndLocs) {
        if (!each.second.empty()) {
            return true;
        }
    }
    return false;
}

/**
 * Form batch from begin iterator parameter of a request buffer
 * If this batch is deemed ready, save the iterator pointing to its last element
 * @param memClk
 * @param begin
 * @return true if this batch is ready, otherwise false
 */
bool SMS::batchFormation(sc_time const &memClk, Thread const &thread,
                         std::deque<gp *> const &requestBuffer, size_t const &inclusiveBeginLoc)
{
    if (requestBuffer.empty()) {
        return false;
    }

    assert(inclusiveBeginLoc <= requestBuffer.size());
    if (requestBuffer.size() == inclusiveBeginLoc) {
        return false;
    }

    if (MPKCs[thread] < LOW_MPKC || isSystemLightlyLoaded()) {
        // bypass requests by forming batch with only one request (threshold age is ZERO)
        readyBatchInclusiveEndLocs[thread].push_back(inclusiveBeginLoc);
        return true;
    } else {
        // forming batch with FIFO size & threshold age constraints
        size_t firstDifferentRowAccessReqLoc = inclusiveBeginLoc;
        Row firstRow = DramExtension::getRow(requestBuffer[inclusiveBeginLoc]);
        bool isBatchReady = false;
        size_t bufferSize = requestBuffer.size();
        while (firstDifferentRowAccessReqLoc != bufferSize
                && DramExtension::getRow(requestBuffer[firstDifferentRowAccessReqLoc]) ==
                firstRow) {
            ++firstDifferentRowAccessReqLoc;
            if (firstDifferentRowAccessReqLoc < bufferSize) {
                if (DramExtension::getRow(requestBuffer[firstDifferentRowAccessReqLoc]) !=
                        firstRow
                        || isExceededReqBufferSize(firstDifferentRowAccessReqLoc)
                        || isThresholdAgeExceeded(thread, memClk, inclusiveBeginLoc,
                                                  firstDifferentRowAccessReqLoc)) {
                    isBatchReady = true;
                    break;
                }
            } else {
                if (isExceededReqBufferSize(firstDifferentRowAccessReqLoc)
                        || isThresholdAgeExceeded(thread, memClk, inclusiveBeginLoc,
                                                  firstDifferentRowAccessReqLoc)) {
                    isBatchReady = true;
                    break;
                }
            }
        }

        // store this ready batch location
        if (isBatchReady) {
            --firstDifferentRowAccessReqLoc;
            readyBatchInclusiveEndLocs[thread].push_back(firstDifferentRowAccessReqLoc);
            debugManager.printDebugMessage(name(),
                                           "Deem batch ready - thread " + to_string(thread.ID()));
            return true;
        } else {
            return false;
        }
    }
}

void SMS::multiBatchFormation(sc_time const &memClk)
{
    for (auto &requestBuffer : requestBuffers) {
        bool formed = false;
        do {
            if (readyBatchInclusiveEndLocs[requestBuffer.first].empty()) {
                formed = batchFormation(memClk, requestBuffer.first, requestBuffer.second, 0);
            } else {
                formed = batchFormation(memClk, requestBuffer.first, requestBuffer.second,
                                        readyBatchInclusiveEndLocs[requestBuffer.first].back() + 1);
            }
        } while (!requestBuffer.second.empty() && formed);
    }
}

gp *SMS::getPendingRequest(Bank bank)
{
    for (const auto &requestBuffer : requestBuffers) {
        for (const auto &request : requestBuffer.second) {
            if (DramExtension::getBank(request) == bank) {
                return request;
            }
        }
    }
    return NULL;
}
