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

#include "MXAnalysisAfterFPGA.h"

#include <span>

#include "spot_finding/DetModuleSpotFinder_cpu.h"
#include "../common/CUDAWrapper.h"
#include "../common/JFJochException.h"
#include "spot_finding/SpotUtils.h"
#include "bragg_prediction/BraggPredictionFactory.h"

double stddev(const std::vector<float> &v) {
    if (v.size() <= 1)
        return 0.0;

    double mean = 0.0f;

    for (const auto &i: v)
        mean += i;
    mean /= v.size();

    double stddev = 0.0f;
    for (const auto &i: v)
        stddev += (i - mean) * (i - mean);

    return sqrt(stddev / (v.size() - 1));
}


MXAnalysisAfterFPGA::MXAnalysisAfterFPGA(const DiffractionExperiment &in_experiment,
                                         const AzimuthalIntegrationMapping &in_integration,
                                         IndexAndRefine &indexer)
    : experiment(in_experiment),
      integration(in_integration),
      indexer(indexer),
      prediction(CreateBraggPrediction(experiment.IsRotationIndexing())) {
    if (experiment.IsSpotFindingEnabled())
        find_spots = true;

    if (experiment.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow())
        cpu_azint = std::make_unique<AzIntEngineCPU>(integration);
}

void MXAnalysisAfterFPGA::RunAzimuthalIntegration(const void *image, AzimuthalIntegrationProfile &profile) {
    if (!cpu_azint)
        return;

    const size_t npixel = experiment.GetPixelsNum();

    auto run = [&](auto tag) {
        using pixel_t = decltype(tag);
        cpu_azint->RunAzint(std::span<const pixel_t>(static_cast<const pixel_t *>(image), npixel), profile);
    };

    switch (experiment.GetByteDepthImage()) {
        case 1:
            experiment.IsPixelSigned() ? run(int8_t{}) : run(uint8_t{});
            break;
        case 2:
            experiment.IsPixelSigned() ? run(int16_t{}) : run(uint16_t{});
            break;
        case 4:
            experiment.IsPixelSigned() ? run(int32_t{}) : run(uint32_t{});
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Pixel depth unsupported");
    }
}

void MXAnalysisAfterFPGA::ReadFromFPGA(const DeviceOutput *output, const SpotFindingSettings &settings, size_t module_number) {
    if (state == State::Disabled || !find_spots || !settings.enable) {
        state = State::Disabled;
    } else {
        const auto t0 = std::chrono::steady_clock::now();
        StrongPixelSet strong_pixel_set;
        strong_pixel_set.ReadFPGAOutput(experiment, *output);
        strong_pixel_set.FindSpots(experiment, settings, spots, module_number);
        const auto t1 = std::chrono::steady_clock::now();
        spot_finding_time_total += (t1 - t0);
        spot_finding_timing_active = true;

        state = State::Enabled;
    }
}

void MXAnalysisAfterFPGA::ReadFromCPU(DeviceOutput *output, const SpotFindingSettings &settings, size_t module_number) {
    std::unique_lock ul(read_from_cpu_mutex);

    if (state == State::Disabled || !find_spots || !settings.enable) {
        state = State::Disabled;
    } else {
        const auto t0 = std::chrono::steady_clock::now();

        state = State::Enabled;

        std::vector<float> d_map(RAW_MODULE_SIZE);

        experiment.CalcSpotFinderResolutionMap(d_map.data(), module_number);

        arr_mean.resize(RAW_MODULE_SIZE);
        arr_sttdev.resize(RAW_MODULE_SIZE);
        arr_valid_count.resize(RAW_MODULE_SIZE);
        arr_strong_pixel.resize(RAW_MODULE_SIZE);

        if (experiment.GetByteDepthImage() == 2)
            FindSpots(*output,
                      settings,
                      d_map.data(),
                      arr_mean.data(),
                      arr_sttdev.data(),
                      arr_valid_count.data(),
                      arr_strong_pixel.data());
        else if (experiment.GetByteDepthImage() == 4)
            FindSpots<int32_t>(*output,
                               settings,
                               d_map.data(),
                               arr_mean.data(),
                               arr_sttdev.data(),
                               arr_valid_count.data(),
                               arr_strong_pixel.data());
        else if (experiment.GetByteDepthImage() == 1)
            FindSpots<int8_t>(*output,
                              settings,
                              d_map.data(),
                              arr_mean.data(),
                              arr_sttdev.data(),
                              arr_valid_count.data(),
                              arr_strong_pixel.data());

        StrongPixelSet strong_pixel_set;
        strong_pixel_set.ReadFPGAOutput(experiment, *output);
        strong_pixel_set.FindSpots(experiment, settings, spots, module_number);

        const auto t1 = std::chrono::steady_clock::now();
        spot_finding_time_total += (t1 - t0);
        spot_finding_timing_active = true;
    }
}

void MXAnalysisAfterFPGA::Process(DataMessage &message, const SpotFindingSettings& spot_finding_settings) {
    if (find_spots && (state == State::Enabled)) {
        const auto t0 = std::chrono::steady_clock::now();
        SpotAnalyze(experiment, spot_finding_settings, spots, message);
        const auto t1 = std::chrono::steady_clock::now();
        spot_finding_time_total += (t1 - t0);

        if (spot_finding_settings.indexing)
            indexer.ProcessImage(message, spot_finding_settings, message.image, *prediction);
    }

    if (spot_finding_timing_active) {
        // total spot-finding time for the whole image
        message.spot_finding_time_s = spot_finding_time_total.count() / 1e6;

        // report/store ms here
        spot_finding_time_total = std::chrono::duration<double, std::micro>{0.0};
        spot_finding_timing_active = false;
    }

    spots.clear();
    state = State::Idle;
}