// 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 "spot_finding/SpotUtils.h"
#include "bragg_prediction/BraggPredictionFactory.h"
#include "image_preprocessing/ImagePreprocessorCPU.h"

#include "azint/AzIntEngineCPU.h"
#include "spot_finding/ImageSpotFinderCPU.h"
#ifdef JFJOCH_USE_CUDA
#include "azint/AzIntEngineGPU.h"
#include "spot_finding/ImageSpotFinderGPU.h"
#include "image_preprocessing/ImagePreprocessorGPU.h"
#include "image_preprocessing/ImagePreprocessorBufferGPU.h"
#include "../common/CUDAWrapper.h"
#endif


MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_experiment,
                                             const AzimuthalIntegration &in_integration,
                                             const PixelMask &in_mask,
                                             IndexAndRefine &in_indexer)
    : experiment(in_experiment),
      integration(in_integration),
      npixels(experiment.GetPixelsNum()),
      xpixels(experiment.GetXPixelsNum()),
      indexer(in_indexer),
      prediction(CreateBraggPrediction(experiment.IsRotationIndexing())),
      mask(in_mask),
      mask_resolution(experiment.GetPixelsNum(), false),
      mask_high_res(-1),
      mask_low_res(-1) {
#ifdef JFJOCH_USE_CUDA
    if (get_gpu_count() == 0) {
#endif
        preprocessor_buffer = std::make_unique<ImagePreprocessorBuffer>(experiment.GetPixelsNum());
        spotFinder = std::make_unique<ImageSpotFinderCPU>(experiment.GetXPixelsNum(), experiment.GetYPixelsNum());
        azint = std::make_unique<AzIntEngineCPU>(integration);
        preprocessor = std::make_unique<ImagePreprocessorCPU>(in_experiment, in_mask);
#ifdef JFJOCH_USE_CUDA
    } else {
        auto stream = std::make_shared<CudaStream>();
        preprocessor_buffer = std::make_unique<ImagePreprocessorBufferGPU>(experiment.GetPixelsNum());
        preprocessor = std::make_unique<ImagePreprocessorGPU>(in_experiment, in_mask, stream);
        spotFinder = std::make_unique<ImageSpotFinderGPU>(experiment.GetXPixelsNum(), experiment.GetYPixelsNum(), stream);
        azint = std::make_unique<AzIntEngineGPU>(integration, stream);
    }
#endif
}

void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
                                    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 auto compression_start_time = std::chrono::steady_clock::now();
    const uint8_t *image_ptr = output.image.GetUncompressedPtr(decompression_buffer);
    const auto compression_end_time = std::chrono::steady_clock::now();
    if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
        output.compression_time_s = std::chrono::duration<float>(compression_end_time - compression_start_time).count();

    const auto preprocessing_start_time = std::chrono::steady_clock::now();
    auto ret = preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
    const auto preprocessing_end_time = std::chrono::steady_clock::now();
    output.preprocessing_time_s = std::chrono::duration<float>(preprocessing_end_time - preprocessing_start_time).count();

    const auto azint_start_time = std::chrono::steady_clock::now();
    azint->Run(*preprocessor_buffer, profile);
    const auto azint_end_time = std::chrono::steady_clock::now();
    output.azint_time_s = std::chrono::duration<float>(azint_end_time - azint_start_time).count();

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

        const auto spot_finding_start_time = std::chrono::steady_clock::now();
        const std::vector<DiffractionSpot> spots = spotFinder->Run(*preprocessor_buffer, spot_finding_settings, mask_resolution);
        SpotAnalyze(experiment, spot_finding_settings, spots, output);
        const auto spot_finding_end_time = std::chrono::steady_clock::now();
        output.spot_finding_time_s = std::chrono::duration<float>(spot_finding_end_time - spot_finding_start_time).count();

        if (spot_finding_settings.indexing)
            indexer.ProcessImage(output, spot_finding_settings,
                CompressedImage(preprocessor_buffer->getBuffer(), experiment.GetXPixelsNum(), experiment.GetYPixelsNum()),
                *prediction);
    }

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

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