// Copyright (2019-2024) Paul Scherrer Institute

#include "MXAnalyzer.h"
#include "CPUSpotFinder.h"
#include "../common/DiffractionGeometry.h"
#include "Regression.h"
#include "../common/CUDAWrapper.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) {
    auto uc = experiment.GetUnitCell();
    if (uc && (get_gpu_count() > 0)) {
        try {
            indexer = std::make_unique<IndexerWrapper>();
            indexer->Setup(uc.value());
        } catch (const std::exception &e) {
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, e.what());
        }
    }

    if (experiment.IsSpotFindingEnabled())
        find_spots = true;
}

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(const int16_t *image, const SpotFindingSettings &settings, size_t module_number) {
    if (!find_spots)
        return;
    if (experiment.GetPixelDepth() == 4)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CPU spot finder simulation doesn't support 32-bit images");
    std::vector<float> d_map(RAW_MODULE_SIZE);

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

    CalcSpotFinderResolutionMap(d_map.data(), experiment, 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);
}

uint32_t MXAnalyzer::FilterSpotsInPowderRings(const std::vector<DiffractionSpot> &spots_filter,
                                              std::vector<DiffractionSpot> &spots_out,
                                              int64_t min_spot_count_ring) {
    uint32_t ret = 0;

    double high_q = 5.0;
    double low_q = 0;
    double q_spacing = 0.01;

    size_t bin_count = (high_q - low_q) / q_spacing + 1;

    std::vector<std::vector<uint32_t> > bins(bin_count);

    for (int i = 0; i < spots_filter.size(); i++) {
        double q = 2 * M_PI / spots_filter[i].GetResolution(experiment);
        if ((q >= low_q) && (q < high_q)) {
            int32_t q_bin = std::floor((q - low_q) / q_spacing);
            bins[q_bin].push_back(i);
        } else // spots outside of azim. int. range are not filtered
            spots_out.push_back(spots_filter[i]);
    }

    for (auto & bin : bins) {
        if (bin.size() > min_spot_count_ring)
            ret += bin.size();
        else {
            for (auto &iter: bin)
                spots_out.push_back(spots_filter[iter]);
        }
    }
    return ret;
}

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

    std::vector<DiffractionSpot> spots_out;
    if (settings.filter_spots_powder_ring) {
        std::vector<DiffractionSpot> spots_no_rings;
        message.spot_count_in_rings = FilterSpotsInPowderRings(spots,
                                                               spots_no_rings,
                                                               settings.min_spot_count_powder_ring);
        FilterSpotsByCount(experiment, spots_no_rings, spots_out);
    } else
        FilterSpotsByCount(experiment, spots, spots_out);

    spots.clear();

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

    if (indexer && settings.indexing) {
        std::vector<Coord> recip;
        recip.reserve(spots_out.size());
        for (const auto &i: spots_out)
            recip.push_back(i.ReciprocalCoord(experiment));

        std::vector<IndexingResult> indexer_result;

        // If there is a really large number of spots detected, it is better to start with a smaller subset
        if (recip.size() > 80)
            indexer_result = indexer->Run(recip, settings.indexing_tolerance, 50);
        if (indexer_result.empty())
            indexer_result = indexer->Run(recip, settings.indexing_tolerance);

        if (!indexer_result.empty()) {
            message.indexing_result = true;
            assert(indexer_result[0].indexed_spot.size() == recip.size());

            // identify indexed spots
            for (int i = 0; i < recip.size(); i++)
                message.spots[i].indexed = indexer_result[0].indexed_spot[i];

            indexer_result[0].l.Save(message.indexing_lattice);
        }
    }
}

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