// Copyright (2019-2023) Paul Scherrer Institute

#include "FPGAAcquisitionDevice.h"
#include <random>
#include <bitset>
#include "../common/to_fixed.h"

void FPGAAcquisitionDevice::StartSendingWorkRequests() {
    stop_work_requests = false;
    send_work_request_future = std::async(std::launch::async, &FPGAAcquisitionDevice::SendWorkRequestThread, this);
}

void FPGAAcquisitionDevice::Finalize() {
    read_work_completion_future.get();

    stop_work_requests = true;
    send_work_request_future.get();

    FPGA_EndAction();

    while (!HW_IsIdle())
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
}

void FPGAAcquisitionDevice::ReadWorkCompletionThread() {
    uint32_t values[16];

    Completion c{};
    bool quit_loop = false;
    do {
        while (!HW_ReadMailbox(values))
            std::this_thread::sleep_for(std::chrono::microseconds(10));

        c = parse_hw_completion(values);
        if (c.data_collection_id == data_collection_id) {
            work_completion_queue.PutBlocking(c);
            if (c.type == Completion::Type::End)
                quit_loop = true;
        } else if (logger) {
            if (c.type == Completion::Type::Start)
                logger->Warning("Stream {} Start completion with wrong data collection ID", data_stream);
            else
                logger->Warning("Stream {} Image completion with wrong data collection ID frame {} module {}",
                                data_stream, c.frame_number, c.module_number);
        }

    } while (!quit_loop);
}

void FPGAAcquisitionDevice::SendWorkRequestThread() {
    while (!stop_work_requests) {
        WorkRequest wr{};
        if (work_request_queue.Get(wr)) {
            if ( !HW_SendWorkRequest(wr.handle)) {
                work_request_queue.Put(wr);
                std::this_thread::sleep_for(std::chrono::microseconds(10));
            }
        }  else {
            std::this_thread::sleep_for(std::chrono::microseconds(10));
        }
    }
}

void FPGAAcquisitionDevice::InitializeIntegrationMap(const DiffractionExperiment &experiment,
                                                     const std::vector<uint16_t> &v) {
    std::vector<float> weights(experiment.GetModulesNum() * RAW_MODULE_SIZE, 1.0);
    InitializeIntegrationMap(experiment, v, weights);
}

void FPGAAcquisitionDevice::InitializeIntegrationMap(const DiffractionExperiment &experiment,
                                                     const std::vector<uint16_t> &v,
                                                     const std::vector<float> &weights) {
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);

    if (v.size() != experiment.GetModulesNum() * RAW_MODULE_SIZE)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch regarding integration map array");

    if (weights.size() != experiment.GetModulesNum() * RAW_MODULE_SIZE)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch regarding weights array");

    size_t modules = experiment.GetModulesNum(data_stream);

    if (modules > 2 * buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Not enough host/FPGA buffers to load all integration map values");

    for (int m = 0; m < modules; m++)
        memcpy(buffer_device[m], v.data() + (offset + m) * RAW_MODULE_SIZE, RAW_MODULE_SIZE * sizeof(uint16_t));


    for (int m = 0; m < modules; m++) {
        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            buffer_device[modules + m][i] = to_fixed(weights[(offset + m) * RAW_MODULE_SIZE + i], 15);
        }
    }
    HW_LoadIntegrationMap(modules);
}

void FPGAAcquisitionDevice::SetInternalGeneratorFrameForAllModules(const std::vector<uint16_t> &v) {
    if (v.size() != RAW_MODULE_SIZE) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Error in size of custom internal generator frame");
    }

    for (int m = 0; m < max_modules; m++) {
        memcpy(internal_pkt_gen_frame.data() + m * RAW_MODULE_SIZE,
               v.data(), RAW_MODULE_SIZE * sizeof(uint16_t));
    }

    if (max_modules > buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Not enough host/FPGA buffers to load all integration map values");

    for (int m = 0; m < max_modules; m++)
        memcpy(buffer_device[m], internal_pkt_gen_frame.data() + m * RAW_MODULE_SIZE, RAW_MODULE_SIZE * sizeof(uint16_t));

    HW_LoadInternalGeneratorFrame(max_modules);
}


void FPGAAcquisitionDevice::SetInternalGeneratorFrame(const std::vector<uint16_t> &v) {
    if (v.empty() || (v.size() % RAW_MODULE_SIZE != 0)) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Error in size of custom internal generator frame");
    }

    size_t nmodules = v.size() / RAW_MODULE_SIZE;
    if (nmodules > max_modules) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Max number of modules exceeded");
    }

    memcpy(internal_pkt_gen_frame.data(), v.data(), nmodules * RAW_MODULE_SIZE * sizeof(uint16_t));

    if (nmodules > buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Not enough host/FPGA buffers to load all integration map values");

    for (int m = 0; m < nmodules; m++)
        memcpy(buffer_device[m], internal_pkt_gen_frame.data() + m * RAW_MODULE_SIZE, RAW_MODULE_SIZE * sizeof(uint16_t));

    HW_LoadInternalGeneratorFrame(nmodules);
}

void FPGAAcquisitionDevice::InitializeCalibration(const DiffractionExperiment &experiment, const JFCalibration &calib) {
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);

    if (calib.GetModulesNum() != experiment.GetModulesNum())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch regarding module count in calibration and experiment description");

    if (calib.GetStorageCellNum() != experiment.GetStorageCellNumber())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch regarding storage cell count in calibration and experiment description");

    size_t modules = experiment.GetModulesNum(data_stream);
    size_t storage_cells = experiment.GetStorageCellNumber();

    if (modules * (3 + 3 * storage_cells) > buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Not enough host/FPGA buffers to load all calibration constants");

    if (modules * (3 + 3 * storage_cells) > LOAD_CALIBRATION_BRAM_SIZE)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "FPGA is not compatible with that many calibration constants");


    for (int m = 0; m < modules; m++) {
        calib.GainCalibration(m).ExportG0(buffer_device[m]);
        calib.GainCalibration(m).ExportG1(buffer_device[m + modules]);
        calib.GainCalibration(m).ExportG2(buffer_device[m + modules * 2]);
    }

    for (int s = 0; s < storage_cells; s++) {
        auto mask = calib.CalculateMask(experiment, s);
        for (int m = 0; m < modules; m++) {
            auto pedestal_g0 = calib.Pedestal(offset + m, 0, s).GetPedestal();
            auto pedestal_g1 = calib.Pedestal(offset + m, 1, s).GetPedestal();
            auto pedestal_g2 = calib.Pedestal(offset + m, 2, s).GetPedestal();
            for (int i = 0; i < RAW_MODULE_SIZE; i++) {
                if (experiment.GetApplyPixelMaskInFPGA() && (mask[(offset + m) * RAW_MODULE_SIZE + i] != 0)) {
                    buffer_device[(3 + 0 * storage_cells + s) * modules + m][i] = 16384;
                    buffer_device[(3 + 1 * storage_cells + s) * modules + m][i] = 16384;
                    buffer_device[(3 + 2 * storage_cells + s) * modules + m][i] = 16384;
                } else {
                    buffer_device[(3 + 0 * storage_cells + s) * modules + m][i] = pedestal_g0[i];
                    buffer_device[(3 + 1 * storage_cells + s) * modules + m][i] = pedestal_g1[i];
                    buffer_device[(3 + 2 * storage_cells + s) * modules + m][i] = pedestal_g2[i];
                }
            }

        }
    }
    HW_LoadCalibration(modules, storage_cells);
}


void FPGAAcquisitionDevice::FillActionRegister(const DiffractionExperiment& x, DataCollectionConfig &job) {
    std::random_device rd;
    std::uniform_int_distribution<uint16_t> dist;
    data_collection_id = dist(rd);

    job.nmodules            = x.GetModulesNum(data_stream) - 1;
    job.nframes             = x.GetFrameNum();
    job.one_over_energy     = std::lround((1<<20)/ x.GetPhotonEnergy_keV());
    job.nstorage_cells      = x.GetStorageCellNumber() - 1;
    job.mode                = data_collection_id << 16;
    job.nsummation          = x.GetSummation() - 1;

    if (x.GetDetectorMode() == DetectorMode::Conversion)
        job.mode |= MODE_CONV;
    if (!x.IsPixelSigned())
        job.mode |= MODE_UNSIGNED;
    if (x.GetPixelDepth() == 4)
        job.mode |= MODE_32BIT;
}

void FPGAAcquisitionDevice::Start(const DiffractionExperiment &experiment, uint32_t flag) {
    if (!HW_IsIdle())
        throw(JFJochException(JFJochExceptionCategory::AcquisitionDeviceError,
                              "Hardware action running prior to start of data acquisition"));
    DataCollectionConfig cfg_in{}, cfg_out{};

    FillActionRegister(experiment, cfg_in);
    cfg_in.mode |= flag;

    HW_WriteActionRegister(&cfg_in);
    HW_ReadActionRegister(&cfg_out);

    if (cfg_out.mode != cfg_in.mode)
        throw JFJochException(JFJochExceptionCategory::AcquisitionDeviceError,
                              "Mismatch between expected and actual values of configuration registers (mode)");
    if (cfg_out.nframes != cfg_in.nframes)
        throw JFJochException(JFJochExceptionCategory::AcquisitionDeviceError,
                              "Mismatch between expected and actual values of configuration registers (Frames per trigger)");
    if (cfg_out.nmodules != cfg_in.nmodules)
        throw JFJochException(JFJochExceptionCategory::AcquisitionDeviceError,
                              "Mismatch between expected and actual values of configuration registers (#modules)");

    FPGA_StartAction(experiment);

    read_work_completion_future = std::async(std::launch::async, &FPGAAcquisitionDevice::ReadWorkCompletionThread, this);
}

std::vector<uint16_t> FPGAAcquisitionDevice::GetInternalGeneratorFrame() const {
    return internal_pkt_gen_frame;
}

void FPGAAcquisitionDevice::SetInternalGeneratorFrame() {
    std::vector<uint16_t> tmp(RAW_MODULE_SIZE);
    for (int i = 0; i < RAW_MODULE_SIZE; i++)
        tmp[i] = i % 65536;
    SetInternalGeneratorFrameForAllModules(tmp);
}

FPGAAcquisitionDevice::FPGAAcquisitionDevice(uint16_t data_stream)
: AcquisitionDevice(data_stream),
internal_pkt_gen_frame(RAW_MODULE_SIZE * MAX_MODULES_FPGA)  {
}

void FPGAAcquisitionDevice::SetSpotFinderParameters(int16_t count_threshold, double snr_threshold) {
    if (snr_threshold < 0) {
        if (logger)
            logger->Warning("Trying to set SNR threshold below zero: {}", snr_threshold );
        snr_threshold = 0;
    } else if (snr_threshold > 64) {
        if (logger)
            logger->Warning("Trying to set SNR threshold too high: {}", snr_threshold );
        snr_threshold = 64;
    }
    SpotFinderParameters params{.count_threshold = count_threshold, .snr_threshold = to_fixed(snr_threshold, 2)};
    HW_SetSpotFinderParameters(params);
}