// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <random>
#include "FPGAAcquisitionDevice.h"
#include "../common/DiffractionGeometry.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;

    if (logger)
        logger->Debug("Starting read work completion thread for stream {}", data_stream);
    Completion c{};
    bool quit_loop = false;
    do {
        bool read;
        try {
            read = HW_ReadMailbox(&values);
        } catch (const JFJochException& e) {
            if (logger)
                logger->ErrorException(e);
            read = false;
        }

        if (read) {
            c = parse_hw_completion(values);
            if (logger) {
                logger->Debug("Stream {} Completion data collection ID {} type {} frame {} module {}",
                              data_stream,
                              c.data_collection_id,
                              (int)c.type,
                              c.frame_number,
                              c.module_number);
            }
            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);
            }
        } else
            std::this_thread::sleep_for(std::chrono::microseconds(10));
    } while (!quit_loop);

    if (logger)
        logger->Debug("Done read work completion thread for stream {}", data_stream);
}

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::InitializeSpotFinderResolutionMap(const float *data, size_t module_number) {
    memcpy(buffer_device[0]->pixels, data, RAW_MODULE_SIZE * sizeof(float));
    buffer_device[0]->module_statistics.module_number = module_number;
    buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_SPOT_FINDER_RES_MAP;
    LoadCalibration(0);
}

void FPGAAcquisitionDevice::InitializeIntegrationMap(const DiffractionExperiment &experiment,
                                                     const std::vector<uint16_t> &v,
                                                     const std::vector<float> &weights) {
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);
    size_t modules = experiment.GetModulesNum(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");

    for (uint32_t m = 0; m < modules; m++) {
        InitializeIntegrationMap(v.data() + (offset + m) * RAW_MODULE_SIZE,
                                 weights.data() + (offset + m) * RAW_MODULE_SIZE,
                                 m);
    }
}

void FPGAAcquisitionDevice::InitializeIntegrationMap(const uint16_t *map, const float *weights, size_t module_number) {
    memcpy(buffer_device[0]->pixels, map, RAW_MODULE_SIZE * sizeof(uint16_t));
    buffer_device[0]->module_statistics.module_number = module_number;
    buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_INTEGRATION_MAP;
    LoadCalibration(0);

    memcpy(buffer_device[0]->pixels, weights, RAW_MODULE_SIZE * sizeof(float));
    buffer_device[0]->module_statistics.module_number = module_number;
    buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_INTEGRATION_WEIGHTS;
    LoadCalibration(0);
}

void FPGAAcquisitionDevice::SetInternalGeneratorFrame(const uint16_t *input, size_t module_number) {
    memcpy(buffer_device[0], input, RAW_MODULE_SIZE * sizeof(uint16_t));
    buffer_device[0]->module_statistics.module_number = module_number;
    buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_FRAME_GEN;
    LoadCalibration(0);
}

void FPGAAcquisitionDevice::InitializeROIMap(const uint16_t *map, size_t module_number) {
    memcpy(buffer_device[0]->pixels, map, RAW_MODULE_SIZE * sizeof(uint16_t));
    buffer_device[0]->module_statistics.module_number = module_number;
    buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_ROI_CALC;
    LoadCalibration(0);
}

void FPGAAcquisitionDevice::InitializePixelMask(const uint32_t *module_mask, size_t module_number) {
    memcpy(buffer_device[0]->pixels, module_mask, RAW_MODULE_SIZE * sizeof(uint32_t));
    buffer_device[0]->module_statistics.module_number = module_number;
    buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_PXL_MASK;
    LoadCalibration(0);
}


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();

    for (uint32_t m = 0; m < modules; m++) {
        if (experiment.IsFixedGainG1()) {
            calib.GainCalibration(m).ExportFixedG1(buffer_device[0]);
            buffer_device[0]->module_statistics.module_number = m;
            buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_GAIN_G1;
            LoadCalibration(0);
        } else {
            if (experiment.IsUsingGainHG0())
                calib.GainCalibration(m).ExportHG0(buffer_device[0]);
            else
                calib.GainCalibration(m).ExportG0(buffer_device[0]);
            buffer_device[0]->module_statistics.module_number = m;
            buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_GAIN_G0;
            LoadCalibration(0);

            calib.GainCalibration(m).ExportG1(buffer_device[0]);
            buffer_device[0]->module_statistics.module_number = m;
            buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_GAIN_G1;
            LoadCalibration(0);

            calib.GainCalibration(m).ExportG2(buffer_device[0]);
            buffer_device[0]->module_statistics.module_number = m;
            buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_GAIN_G2;
            LoadCalibration(0);
        }
    }

    for (int s = 0; s < storage_cells; s++) {
        for (int m = 0; m < modules; m++) {
            auto pedestal_g1 = calib.Pedestal(offset + m, 1, s).GetPedestal();
            memcpy(buffer_device[0]->pixels, pedestal_g1, RAW_MODULE_SIZE * sizeof(uint16_t));
            buffer_device[0]->module_statistics.module_number = m + modules * s;
            buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_PEDESTAL_G1;
            LoadCalibration(0);

            if (!experiment.IsFixedGainG1()) {
                auto pedestal_g0 = calib.Pedestal(offset + m, 0, s).GetPedestal();
                memcpy(buffer_device[0]->pixels, pedestal_g0, RAW_MODULE_SIZE * sizeof(uint16_t));
                buffer_device[0]->module_statistics.module_number = m + modules * s;
                buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_PEDESTAL_G0;
                LoadCalibration(0);

                auto pedestal_g2 = calib.Pedestal(offset + m, 2, s).GetPedestal();
                memcpy(buffer_device[0]->pixels, pedestal_g2, RAW_MODULE_SIZE * sizeof(uint16_t));
                buffer_device[0]->module_statistics.module_number = m + modules * s;
                buffer_device[0]->module_statistics.load_calibration_destination = LOAD_CALIBRATION_DEST_PEDESTAL_G2;
                LoadCalibration(0);
            }
        }
    }
}

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

    // Avoid data_collection_id of DATA_COLLECTION_ID_PURGE  - this is reserved for cleaning unused work requests from the FPGA
    data_collection_id = dist(rd);
    while (data_collection_id == DATA_COLLECTION_ID_PURGE)
        data_collection_id = dist(rd);

    job.nmodules            = x.GetModulesNum(data_stream) - 1;
    job.nframes             = x.GetFrameNum();
    job.energy_kev          = x.GetPhotonEnergyForConversion_keV();
    job.nstorage_cells      = x.GetStorageCellNumber() - 1;
    job.mode                = data_collection_id << 16;
    job.nsummation          = x.GetFPGASummation() - 1;
    job.data_stream         = data_stream;

    expected_descriptors_per_module = DMA_DESCRIPTORS_PER_MODULE;

    if (x.IsJungfrauConvPhotonCnt())
        job.mode |= MODE_CONV;

    if (x.IsFixedGainG1())
        job.mode |= MODE_FIXG1;

    if (!x.IsPixelSigned())
        job.mode |= MODE_UNSIGNED;

    if (x.GetEigerBitDepth() == 32)
        job.mode |= MODE_EIGER_32BIT;
    else if (x.GetEigerBitDepth() == 8)
        job.mode |= MODE_EIGER_8BIT;

    if (x.GetByteDepthFPGA() == 4)
        job.mode |= MODE_32BIT_OUTPUT;
    if (x.GetByteDepthFPGA() == 1)
        job.mode |= MODE_8BIT_OUTPUT;

    if (x.IsApplyPixelMask())
        job.mode |= MODE_APPLY_PIXEL_MASK;

    if (x.GetLossyCompressionPoisson()) {
        job.mode |= MODE_SQROOT;
        job.sqrtmult = x.GetLossyCompressionPoisson().value();
    }

    if (x.GetPixelValueLowThreshold()) {
        job.mode |= MODE_THRESHOLD;
        job.pxlthreshold = x.GetPixelValueLowThreshold().value();
    }
}

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);
}

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

uint32_t FPGAAcquisitionDevice::GetExpectedDescriptorsPerModule() const {
    return expected_descriptors_per_module;
}

void FPGAAcquisitionDevice::LoadCalibration(uint32_t handle) {
    HW_LoadCalibration(LoadCalibrationConfig{.handle = handle});
}