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

#include "MXAnalysisWithoutFPGA.h"

#include "spot_finding/StrongPixelSet.h"
#include "../compression/JFJochDecompress.h"

#include "SpotAnalyze.h"
#include "spot_finding/ImageSpotFinderFactory.h"
#include "bragg_integration/BraggPredictionFactory.h"

MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_experiment,
                                             const AzimuthalIntegration &in_integration,
                                             const PixelMask &in_mask,
                                             IndexerThreadPool *in_indexer)
    : experiment(in_experiment),
      integration(in_integration),
      roi_map(experiment.ExportROIMap()),
      roi_names(experiment.ROI().GetROINameMap()),
      roi_count(experiment.ROI().size()),
      npixels(experiment.GetPixelsNum()),
      xpixels(experiment.GetXPixelsNum()),
      mask_1bit(npixels, false),
      spotFinder(CreateImageSpotFinder(experiment.GetXPixelsNum(), experiment.GetYPixelsNum())),
      indexer(in_indexer),
      prediction(CreateBraggPrediction()),
      updated_image(spotFinder->GetInputBuffer()),
      azint_bins(in_integration.GetBinNumber()),
      saturation_limit(experiment.GetSaturationLimit()),
      mask(in_mask),
      mask_resolution(experiment.GetPixelsNum(), false),
      mask_high_res(-1),
      mask_low_res(-1) {
    for (int i = 0; i < npixels; i++)
        mask_1bit[i] = (in_mask.GetMask().at(i) != 0);
}

void MXAnalysisWithoutFPGA::Analyze(DataMessage &output, std::vector<uint8_t> &image,
                                    AzimuthalIntegrationProfile &profile,
                                    const SpotFindingSettings &spot_finding_settings) {
    if ((output.image.GetWidth() != xpixels)
        || (output.image.GetWidth() * output.image.GetHeight() != npixels))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Mismatch in pixel size");

    const uint8_t *image_ptr = output.image.GetUncompressedPtr(image);

    switch (output.image.GetMode()) {
        case CompressedImageMode::Int8:
            Analyze<int8_t>(output, image_ptr, INT8_MIN, INT8_MAX, profile, spot_finding_settings);
            break;
        case CompressedImageMode::Int16:
            Analyze<int16_t>(output, image_ptr, INT16_MIN, INT16_MAX, profile, spot_finding_settings);
            break;
        case CompressedImageMode::Int32:
            Analyze<int32_t>(output, image_ptr, INT32_MIN, INT32_MAX, profile, spot_finding_settings);
            break;
        case CompressedImageMode::Uint8:
            Analyze<uint8_t>(output, image_ptr, UINT8_MAX, UINT8_MAX, profile, spot_finding_settings);
            break;
        case CompressedImageMode::Uint16:
            Analyze<uint16_t>(output, image_ptr, UINT16_MAX, UINT16_MAX, profile, spot_finding_settings);
            break;
        case CompressedImageMode::Uint32:
            Analyze<uint32_t>(output, image_ptr, UINT32_MAX, UINT32_MAX, profile, spot_finding_settings);
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "RGB/float mode not supported");
    }
}

void MXAnalysisWithoutFPGA::UpdateMaskResolution(const SpotFindingSettings &settings) {
    mask_low_res = settings.low_resolution_limit;
    mask_high_res = settings.high_resolution_limit;
    auto const &resolution_map = integration.Resolution();
    for (int i = 0; i < mask_resolution.size(); i++)
        mask_resolution[i] = (resolution_map[i] > mask_low_res) || (resolution_map[i] < mask_high_res);
}

template<class T>
void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
                                    const uint8_t *in_image,
                                    T err_pixel_val,
                                    T sat_pixel_val,
                                    AzimuthalIntegrationProfile &profile,
                                    const SpotFindingSettings &settings) {
    auto image = reinterpret_cast<const T *>(in_image);

    std::vector<ROIMessage> roi(roi_count);

    std::vector<float> azim_sum(azint_bins, 0.0f);
    //std::vector<float> azim_sum2(integration.GetBinNumber(), 0.0f);
    std::vector<uint32_t> azim_count(azint_bins, 0);

    size_t err_pixels = 0;
    size_t masked_pixels = 0;
    size_t sat_pixels = 0;
    int64_t min_value = INT64_MAX;
    int64_t max_value = INT64_MIN;

    if (sat_pixel_val > saturation_limit)
        sat_pixel_val = static_cast<T>(saturation_limit);

    auto &pixel_to_bin = integration.GetPixelToBin();
    auto &corrections = integration.Corrections();

    profile.Clear(integration);

    for (int i = 0; i < npixels; i++) {
        if (mask_1bit[i] != 0) {
            updated_image[i] = INT32_MIN;
            ++masked_pixels;
        } else if (image[i] >= sat_pixel_val) {
            updated_image[i] = INT32_MIN;
            ++sat_pixels;
        } else if (std::is_signed<T>::value && (image[i] == err_pixel_val)) {
            // Error pixels are possible only for signed types
            updated_image[i] = INT32_MIN;
            ++err_pixels;
        } else {
            updated_image[i] = static_cast<int32_t>(image[i]);

            if (image[i] > max_value)
                max_value = image[i];
            if (image[i] < min_value)
                min_value = image[i];

            if (roi_count > 0 && (roi_map[i] != 0)) {
                int64_t x = i % xpixels;
                int64_t y = i / xpixels;

                for (int8_t r = 0; r < roi_count; r++) {
                    if ((roi_map[i] & (1 << r)) != 0) {
                        roi[r].sum += image[i];
                        roi[r].sum_square += image[i] * image[i];
                        roi[r].pixels += 1;
                        if (image[i] > roi[r].max_count)
                            roi[r].max_count = image[i];
                        roi[r].x_weighted += x * image[i];
                        roi[r].y_weighted += y * image[i];
                    }
                }
            }
            const uint16_t bin = pixel_to_bin[i];
            if (bin < azint_bins) {
                float val = image[i] * corrections[i];
                azim_sum[bin] += val;
                //azim_sum2[bin] += val * val;
                ++azim_count[bin];
            }
        }
    }

    if (settings.enable) {
        // Update resolution mask
        if (mask_high_res != settings.high_resolution_limit
            || mask_low_res != settings.low_resolution_limit)
            UpdateMaskResolution(settings);

        const std::vector<DiffractionSpot> spots = spotFinder->Run(settings, mask_resolution);

        SpotAnalyze(experiment, settings, spots,
                    CompressedImage(updated_image, experiment.GetXPixelsNum(), experiment.GetYPixelsNum()),
                    *prediction, indexer, output);
    }

    profile.Add(azim_sum, azim_count);

    output.max_viable_pixel_value = max_value;
    output.min_viable_pixel_value = min_value;
    output.error_pixel_count = err_pixels;
    output.saturated_pixel_count = sat_pixels;
    output.az_int_profile = profile.GetResult();
    output.bkg_estimate = profile.GetBkgEstimate(integration.Settings());

    for (const auto &[key, val]: roi_names)
        output.roi[key] = roi[val];
}