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

#include "MXAnalysisWithoutFPGA.h"

#include <algorithm>

#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 "roi/ROIIntegrationCPU.h"
#include "spot_finding/ImageSpotFinderCPU.h"
#include "spot_finding/AdaptiveSpotFinderCPU.h"
#include "bragg_integration/BraggIntegrationEngineCPU.h"
#ifdef JFJOCH_USE_CUDA
#include "azint/AzIntEngineGPU.h"
#include "roi/ROIIntegrationGPU.h"
#include "spot_finding/ImageSpotFinderGPU.h"
#include "spot_finding/AdaptiveSpotFinderGPU.h"
#include "image_preprocessing/ImagePreprocessorGPU.h"
#include "image_preprocessing/ImagePreprocessorBufferGPU.h"
#include "bragg_integration/BraggIntegrationEngineGPU.h"
#include "../common/CUDAWrapper.h"
#endif


MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_experiment,
                                             const AzimuthalIntegrationMapping &in_integration,
                                             const PixelMask &in_mask,
                                             IndexAndRefine &in_indexer,
                                             bool in_enable_fused_adaptive_gpu)
    : experiment(in_experiment),
      integration(in_integration),
      enable_fused_adaptive_gpu(in_enable_fused_adaptive_gpu),
      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);
        bragg_engine = std::make_unique<BraggIntegrationEngineCPU>(in_experiment);
        if (experiment.ROI().size() >= 1)
            roi = std::make_unique<ROIIntegrationCPU>(experiment);
#ifdef JFJOCH_USE_CUDA
    } else {
        stream = std::make_shared<CudaStream>();
        preprocessor_buffer = std::make_unique<ImagePreprocessorBufferGPU>(experiment.GetPixelsNum());
        // The preprocessed image only has to come back to the host if a CPU engine reads it. Every
        // engine built below runs on the GPU, except the CPU adaptive finder that is kept when the fused
        // GPU engine is off (the online receiver) - so that is the one case that needs the copy.
        preprocessor = std::make_unique<ImagePreprocessorGPU>(in_experiment, in_mask, stream,
                                                             /*copy_image_to_host=*/!enable_fused_adaptive_gpu);
        spotFinder = std::make_unique<ImageSpotFinderGPU>(experiment.GetXPixelsNum(), experiment.GetYPixelsNum(), stream);
        azint = std::make_unique<AzIntEngineGPU>(integration, stream);
        bragg_engine = std::make_unique<BraggIntegrationEngineGPU>(in_experiment, stream);
        if (experiment.ROI().size() >= 1)
            roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
        if (enable_fused_adaptive_gpu) {
            // One GPU engine that computes the azimuthal profile and the adaptive spot mask in a single
            // image pass. fused_adaptive aliases it so Analyze() can lift the profile out of it.
            auto fused = std::make_unique<AdaptiveSpotFinderGPU>(integration, stream);
            fused_adaptive = fused.get();
            adaptiveSpotFinder = std::move(fused);
        }
    }
#endif
    if (!adaptiveSpotFinder)
        adaptiveSpotFinder = std::make_unique<AdaptiveSpotFinderCPU>(integration);
}

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 = Decompress(output.image);
    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();

    // The fused GPU engine (rugnux offline, GPU, adaptive detection) produces the azimuthal profile as
    // a byproduct of spot finding, so the separate azint pass is skipped in that case and the profile is
    // lifted out of the finder below.
    const bool fused = enable_fused_adaptive_gpu && spot_finding_settings.enable
                    && spot_finding_settings.adaptive_threshold && fused_adaptive != nullptr;

    if (!fused) {
        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 (roi)
        roi->Run(*preprocessor_buffer, output.roi);

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

        ImageSpotFinder &finder = spot_finding_settings.adaptive_threshold
                                      ? static_cast<ImageSpotFinder &>(*adaptiveSpotFinder)
                                      : *spotFinder;
        const auto integrate_fn = [this](const std::vector<Reflection> &predicted, size_t npredicted,
                                         int64_t image_number) {
            return bragg_engine->Run(*preprocessor_buffer, predicted, npredicted, image_number);
        };

        // A missing min-pix (std::nullopt) means "choose it per image". This applies only to the stills
        // indexing path (each frame is indexed independently); rotation indexing builds one lattice from
        // all frames, so it keeps the fixed min-pix and the single-pass finder.
        const bool adaptive_min_pix = !spot_finding_settings.min_pix_per_spot.has_value()
                                   && spot_finding_settings.indexing
                                   && !experiment.IsRotationIndexing();
        if (adaptive_min_pix) {
            // Choose the per-image min-pix adaptively instead of a fixed one. min-pix filters connected
            // components AFTER detection, so BOTH the detection (the expensive per-pixel pass) and the
            // connected-component search run ONCE and only the filter is repeated; the azimuthal
            // profile is the one Detect() computed.
            // Index at 3/2/1 (index-only, no integration/accumulation), keep whichever maximises
            // n_indexed^2 / n_total (indexed count weighted by indexed fraction) together with its spot
            // list, and integrate that one.
            const auto detect_start_time = std::chrono::steady_clock::now();
            finder.Detect(*preprocessor_buffer, spot_finding_settings);
            const auto &components = finder.ExtractComponents(*preprocessor_buffer, spot_finding_settings);
            float spot_finding_time_s =
                    std::chrono::duration<float>(std::chrono::steady_clock::now() - detect_start_time).count();
            float indexing_time_s = 0.0f;

            SpotFindingSettings s = spot_finding_settings;
            std::vector<DiffractionSpot> best_spots;
            int best_mp = 0;
            double best_score = -1.0;
            for (int mp : {3, 2, 1}) {
                s.min_pix_per_spot = mp;
                const auto extract_start_time = std::chrono::steady_clock::now();
                std::vector<DiffractionSpot> spots = ImageSpotFinder::Filter(components, s);
                spot_finding_time_s +=
                        std::chrono::duration<float>(std::chrono::steady_clock::now() - extract_start_time).count();
                SpotAnalyze(experiment, s, spots, output);
                const bool indexed = indexer.IndexFrameOnly(output, s);
                indexing_time_s += output.indexing_time_s.value_or(0.0f);
                if (indexed) {
                    const double n_idx = static_cast<double>(output.spot_count_indexed.value_or(0));
                    const double n_tot = static_cast<double>(std::max<int64_t>(1, output.spot_count.value_or(1)));
                    const double score = n_idx * n_idx / n_tot;
                    if (score > best_score) {
                        best_score = score;
                        best_mp = mp;
                        best_spots = std::move(spots);
                    }
                }
            }
            if (best_mp != 0) {
                // Index and integrate the winning spot list; no spot finding left to do.
                s.min_pix_per_spot = best_mp;
                SpotAnalyze(experiment, s, best_spots, output);
                indexer.ProcessImage(output, s, *prediction, integrate_fn);
                indexing_time_s += output.indexing_time_s.value_or(0.0f);
            }
            // Each indexer call reports only its own time, so the escalation's total is summed here.
            output.spot_finding_time_s = spot_finding_time_s;
            output.indexing_time_s = indexing_time_s;
        } else {
            const auto spot_finding_start_time = std::chrono::steady_clock::now();
            const std::vector<DiffractionSpot> spots = finder.Run(*preprocessor_buffer, spot_finding_settings);
            SpotAnalyze(experiment, spot_finding_settings, spots, output);
            output.spot_finding_time_s = std::chrono::duration<float>(std::chrono::steady_clock::now() - spot_finding_start_time).count();
            if (spot_finding_settings.indexing)
                indexer.ProcessImage(output, spot_finding_settings, *prediction, integrate_fn);
        }

#ifdef JFJOCH_USE_CUDA
        if (fused) {
            // Lift the azimuthal profile the fused engine computed in the same detection pass; its azint
            // cost is folded into spot_finding_time_s above.
            profile.Clear(integration);
            profile += fused_adaptive->GetProfile();
            output.azint_time_s = 0.0f;
        }
#endif
    }

    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.az_int_profile_count = profile.GetPixelCount();
    output.az_int_profile_std = profile.GetStd();

    output.bkg_estimate = profile.GetBkgEstimate(integration.Settings());
    output.ice_ring_score = profile.GetIceRingScore(integration.Settings(),
                                                    spot_finding_settings.ice_ring_width_Q_recipA);
}

void MXAnalysisWithoutFPGA::RebuildROI() {
    if (experiment.ROI().empty()) {
        roi.reset();
        return;
    }
#ifdef JFJOCH_USE_CUDA
    if (stream) {
        roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
        return;
    }
#endif
    roi = std::make_unique<ROIIntegrationCPU>(experiment);
}

void MXAnalysisWithoutFPGA::AnalyzeROIOnly(DataMessage &output) {
    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 = Decompress(output.image);
    preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
    RunROIOnly(output);
}

const uint8_t *MXAnalysisWithoutFPGA::Decompress(const CompressedImage &image) {
    // An uncompressed image is read straight out of the message and never touches decompression_buffer,
    // so it stays in pageable memory - the buffer is only worth page-locking when it is actually used.
    if (image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
        preprocessor->PinInputBuffer(decompression_buffer, image.GetUncompressedSize());
    return image.GetUncompressedPtr(decompression_buffer);
}

void MXAnalysisWithoutFPGA::RunROIOnly(DataMessage &output) {
    output.roi.clear();
    if (roi)
        roi->Run(*preprocessor_buffer, output.roi);
}

void MXAnalysisWithoutFPGA::UpdateMaskResolution(const SpotFindingSettings &settings) {
    mask_low_res = settings.low_resolution_limit;
    mask_high_res = settings.high_resolution_limit;
    // No high-resolution limit requested -> mask nothing at the high-resolution end: no pixel has d < 0,
    // and the detector's own edge is where the pixels stop anyway.
    const float high_res = mask_high_res.value_or(0.0f);
    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] < high_res);

    // The finders keep their own copy (the GPU ones a bit-packed device copy), so the mask is handed
    // over here - when the limits change - rather than with every image.
    spotFinder->SetResolutionMask(mask_resolution);
    adaptiveSpotFinder->SetResolutionMask(mask_resolution);
}