// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include <thread>

#include "AcquisitionCounters.h"
#include "../common/JFJochException.h"

AcquisitionCounters::AcquisitionCounters()
        : head(max_modules, 0), slowest_head(0), total_packets(0), expected_frames(0), acquisition_finished(false) {}

void AcquisitionCounters::Reset(const DiffractionExperiment &experiment, uint16_t data_stream) {
    std::unique_lock<std::shared_mutex> ul(m);
    acquisition_finished = false;

    slowest_head = 0;

    if ((experiment.GetDetectorMode() == DetectorMode::PedestalG0) ||
        (experiment.GetDetectorMode() == DetectorMode::PedestalG1) ||
        (experiment.GetDetectorMode() == DetectorMode::PedestalG2))
        expected_frames = experiment.GetFrameNum();
    else
        expected_frames = experiment.GetImageNum() * experiment.GetSummation();

    nmodules = experiment.GetModulesNum(data_stream);
    if (nmodules > max_modules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Acquisition counter cannot support that many modules");
    for (int i = 0; i < max_modules; i++)
        head[i] = 0;

    handle_for_frame = std::vector<uint64_t>((expected_frames+1) * nmodules, HandleNotFound);
    full_module_collected = std::vector<uint8_t>(expected_frames * nmodules);
    saved_completions = std::vector<Completion>(expected_frames * nmodules);
    packets_per_module = std::vector<uint64_t>(nmodules);
    total_packets = 0;
}

void AcquisitionCounters::UpdateCounters(const Completion *c) {
    std::unique_lock<std::shared_mutex> ul(m);

    if (c->module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "UpdateCounters wrong module number: " + std::to_string(c->module_number) + " for frame " + std::to_string(c->frame_number));
    if (c->frame_number >= expected_frames)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "UpdateCounters frame number is out of bounds");
    else {
        if (head.at(c->module_number) < c->frame_number)
            head.at(c->module_number) = c->frame_number;

        if (c->frame_number > slowest_head) {
            slowest_head = head[0];
            for (int i = 1; i < nmodules; i++)
                if (head[i] < slowest_head)
                    slowest_head = head[i];
        }

        full_module_collected.at(c->frame_number * nmodules + c->module_number) = (c->packet_count == expected_packets_per_module);
        handle_for_frame.at(c->frame_number * nmodules + c->module_number) = c->handle;
        saved_completions.at(c->frame_number * nmodules + c->module_number) = *c;

        total_packets += c->packet_count;
        packets_per_module[c->module_number] += c->packet_count;
    }

    data_updated.notify_all();
}

void AcquisitionCounters::SetAcquisitionFinished() {
    std::unique_lock<std::shared_mutex> ul(m);
    acquisition_finished = true;
    data_updated.notify_all();
}

uint64_t AcquisitionCounters::GetBufferHandle(size_t frame, uint16_t module_number) const {
    if (frame >= expected_frames)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetBufferHandle Wrong frame number: " + std::to_string(frame));
    if (module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetBufferHandle Wrong module number: " + std::to_string(module_number)
                              + " for frame " + std::to_string(frame));
    return handle_for_frame.at(frame * nmodules + module_number);
}

uint64_t AcquisitionCounters::GetBufferHandleAndClear(size_t frame, uint16_t module_number) {
    std::unique_lock<std::shared_mutex> ul(m);
    if (frame >= expected_frames)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetBufferHandleAndClear Wrong frame number: " + std::to_string(frame));
    if (module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetBufferHandleAndClear Wrong module number: " + std::to_string(module_number)
                              + " for frame " + std::to_string(frame));
    uint64_t ret_val = handle_for_frame.at(frame * nmodules + module_number);
    handle_for_frame.at(frame * nmodules + module_number) = HandleNotFound;
    return ret_val;
}

uint64_t AcquisitionCounters::GetHead(uint16_t module_number) const {
    if (module_number >= max_modules)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetHead Wrong module number: " + std::to_string(module_number));
    return head[module_number];
}

uint64_t AcquisitionCounters::GetSlowestHead() const {
    return slowest_head;
}

void AcquisitionCounters::WaitForFrame(size_t curr_frame, uint16_t module_number) const {
    uint64_t slowest_head_tmp = (module_number == UINT16_MAX) ? GetSlowestHead() : GetHead(module_number);
    while (!acquisition_finished && (slowest_head_tmp < curr_frame)) {
        std::this_thread::sleep_for(std::chrono::microseconds(100));
        slowest_head_tmp = (module_number == UINT16_MAX) ? GetSlowestHead() : GetHead(module_number);
    }
}

int64_t AcquisitionCounters::CalculateDelay(size_t curr_frame, uint16_t module_number) const {
    uint64_t slowest_head_tmp;

    if (module_number == UINT16_MAX)
        slowest_head_tmp = GetSlowestHead();
    else
        slowest_head_tmp = GetHead(module_number);

    return slowest_head_tmp - curr_frame;
}

bool AcquisitionCounters::IsAcquisitionFinished() const {
    return acquisition_finished;
}

bool AcquisitionCounters::IsFullModuleCollected(size_t frame, uint16_t module_number) const {
    if (frame >= expected_frames)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "IsFullModuleCollected Wrong frame number: " + std::to_string(frame));

    if (module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "IsFullModuleCollected Wrong module number: " + std::to_string(module_number)
                              + " for frame " + std::to_string(frame));

    return full_module_collected[frame * nmodules + module_number];
}

Completion AcquisitionCounters::GetCompletion(size_t frame, uint16_t module_number) const {
    if (frame >= expected_frames)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetBunchID Wrong frame number: " + std::to_string(frame));

    if (module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetBunchID Wrong module number: " + std::to_string(module_number)
                              + " for frame " + std::to_string(frame));

    return saved_completions[frame * nmodules + module_number];
}

uint64_t AcquisitionCounters::GetTotalPackets() const {
    return total_packets;
}

uint64_t AcquisitionCounters::GetTotalPackets(uint16_t module_number) const {
    std::unique_lock<std::shared_mutex> ul(m);

    if (module_number >= nmodules)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              "GetTotalPackets Wrong module number: " + std::to_string(module_number));\

    return packets_per_module[module_number];
}

uint64_t AcquisitionCounters::GetExpectedPackets() const {
    return GetExpectedPacketsPerModule() * nmodules;
}

uint64_t AcquisitionCounters::GetExpectedPacketsPerModule() const {
    return expected_frames * expected_packets_per_module;
}