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

#include "MXAnalyzer.h"
#include "CPUSpotFinder.h"
#include "../common/CUDAWrapper.h"
#include "indexing/IndexerFactory.h"
#include "SpotAnalysis.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));
}


MXAnalyzer::MXAnalyzer(const DiffractionExperiment &in_experiment)
: experiment(in_experiment) {
    if (experiment.IsSpotFindingEnabled())
        find_spots = true;
}

MXAnalyzer &MXAnalyzer::SetIndexer(IndexerThreadPool *input) {
    indexer = input;
    return *this;
}

void MXAnalyzer::ReadFromFPGA(const DeviceOutput *output, const SpotFindingSettings &settings, size_t module_number) {
    if (!find_spots || !settings.enable)
        return;
    StrongPixelSet strong_pixel_set;
    strong_pixel_set.ReadFPGAOutput(experiment, *output);
    strong_pixel_set.FindSpots(experiment, settings, spots, module_number);
}

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

    if (!find_spots)
        return;

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

    ReadFromFPGA(output, settings, module_number);
}

void MXAnalyzer::ReadFromCPU(const int16_t *image,
                             const SpotFindingSettings &settings,
                             size_t module_number) {
    std::unique_lock ul(read_from_cpu_mutex);

    if (!find_spots)
        return;
    if (experiment.GetByteDepthImage() != 2)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CPU spot finder simulation supports only 16-bit images");
    std::vector<float> d_map(RAW_MODULE_SIZE);

    DeviceOutput output{};
    memcpy(output.pixels, image, RAW_MODULE_SIZE * sizeof(int16_t));

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

    arr_mean.resize(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    arr_sttdev.resize(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    arr_valid_count.resize(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    arr_strong_pixel.resize(experiment.GetModulesNum() * RAW_MODULE_SIZE);

    FindSpots(output,
              settings,
              d_map.data(),
              arr_mean.data() + module_number * RAW_MODULE_SIZE,
              arr_sttdev.data() + module_number * RAW_MODULE_SIZE,
              arr_valid_count.data() + module_number * RAW_MODULE_SIZE,
              arr_strong_pixel.data() + module_number * RAW_MODULE_SIZE);

    ReadFromFPGA(&output, settings, module_number);
}

void MXAnalyzer::Process(DataMessage &message, const SpotFindingSettings& settings) {
    if (!find_spots)
        return;

    std::vector<SpotToSave> spots_out;
    auto geom = experiment.GetDiffractionGeometry();

    for (const auto &spot: spots)
        spots_out.push_back(spot.Export(geom));

    CountSpots(message, spots_out, settings.cutoff_spot_count_low_res);

    spots.clear();

    FilterSpotsByCount(experiment.GetMaxSpotCount(), spots_out, message.spots);

    if (indexer && settings.indexing) {
        auto latt = indexer->Run(experiment, message).get();

        if (latt && settings.quick_integration) {
            auto res = BraggIntegrate2D(experiment, message.image, latt.value());

            message.reflections = res.reflections;
            message.b_factor = res.b_factor;
        }
    }
}

const std::vector<float> &MXAnalyzer::GetCPUMean() const {
    return arr_mean;
}

const std::vector<float> &MXAnalyzer::GetCPUStdDev() const {
    return arr_sttdev;
}

const std::vector<uint32_t> &MXAnalyzer::GetCPUValidCount() const {
    return arr_valid_count;
}

const std::vector<uint32_t> &MXAnalyzer::GetCPUStrongPixel() const {
    return arr_strong_pixel;
}