// 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 <spdlog/spdlog.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 "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 - so that is the one case that needs the copy. Every caller currently passes
        // enable_fused_adaptive_gpu = true, so on the GPU path the copy is off in practice.
        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");

    // Decompress on the device where the preprocessor can, so only the compressed chunk crosses PCIe
    // and the host does no decompression at all. AnalyzeCompressed says whether it took the image;
    // when it declines (a CPU preprocessor, or an algorithm with no device decoder) fall through to
    // the host route unchanged. The two produce the same preprocessed image.
    const auto compression_start_time = std::chrono::steady_clock::now();
    ImageStatistics ret{};
    bool decoded_on_device = false;
    try {
        decoded_on_device = preprocessor->AnalyzeCompressed(*preprocessor_buffer, output.image, ret);
    } catch (const JFJochException &e) {
        // The device route must never be the reason a frame fails: whatever it could not handle, the
        // host decoder gets its turn. If the data really is bad the host throws too and the caller
        // sees the same error it saw before any of this existed - but a GPU-side problem costs speed
        // rather than the acquisition.
        spdlog::warn("Device decoding failed ({}), falling back to host decompression", e.what());
        decoded_on_device = false;
    }
    const auto compression_end_time = std::chrono::steady_clock::now();

    if (!decoded_on_device) {
        const uint8_t *image_ptr = Decompress(output.image);
        const auto decompressed_time = std::chrono::steady_clock::now();
        if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
            output.compression_time_s = std::chrono::duration<float>(decompressed_time - compression_start_time).count();

        const auto preprocessing_start_time = std::chrono::steady_clock::now();
        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();
    } else {
        // Decode and preprocess are one device operation here, but the decompression is still a real,
        // separately measurable cost - the decoder brackets it with CUDA events - so it is still
        // reported as one. Leaving compression_time_s unset instead would blank the broker's
        // "compression" plot trace and fill /entry/profiling/compressionTime with NaN.
        const float total_s = std::chrono::duration<float>(compression_end_time - compression_start_time).count();
        const float decompress_s = std::min(preprocessor->GetLastDecompressionTime_s(), total_s);
        output.compression_time_s = decompress_s;
        output.preprocessing_time_s = total_s - decompress_s;
    }

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

        // The radial background correction has to be decided BEFORE this image is integrated, so the
        // ice score is taken here rather than with the other per-image quantities at the end of the
        // function. It needs the peak-excluded per-ring background, which the adaptive finder has as
        // soon as it has detected - so this must be called after detection and before integration.
        // Where no such background exists (no adaptive finder), auto leaves the correction off: the
        // plain profile carries the Bragg peaks and cannot support an absolute threshold.
        const auto decide_radial_background = [this, &spot_finding_settings, &output]() {
            if (!bragg_engine->IsBackgroundRadialAuto())
                return;
            // Same condition as the score at the end of this function: the adaptive finder holds its
            // ring background from whenever it last ran, so requiring that it ran for THIS image is
            // what keeps a stale curve out.
            if (!spot_finding_settings.adaptive_threshold)
                return;
            const std::vector<float> &ring_bkg = adaptiveSpotFinder->GetRingBackground();
            if (ring_bkg.empty())
                return;
            output.ice_ring_score = AzimuthalIntegrationProfile::IceRingScore(
                    ring_bkg, integration.GetQBinCount(), integration.Settings(),
                    spot_finding_settings.ice_ring_width_Q_recipA);
            bragg_engine->BackgroundRadial(*output.ice_ring_score
                                           >= experiment.GetScalingSettings().GetIceMinScore());
        };

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

    // The ice score wants a radial profile with the Bragg peaks taken OUT of it. The azimuthal profile
    // is a plain per-ring mean, so a strong low-resolution reflection landing in a ring's bin is
    // indistinguishable from ice sitting there - measured, that alone lifts clean crystals to a score
    // of 1.5-4.2, right into the range real ice occupies. The adaptive spot finder already computes
    // exactly what is wanted: a sigma-clipped per-ring background, in the same bins, from which the
    // peaks have been removed (an ice ring is azimuthally smooth, so it survives the clip). It is in
    // raw counts rather than corrected ones, which the score does not care about - it is a ratio to the
    // background interpolated under the ring, and the corrections are smooth in radius.
    const std::vector<float> &ring_bkg = adaptiveSpotFinder->GetRingBackground();
    const bool have_ring_bkg = spot_finding_settings.enable && spot_finding_settings.adaptive_threshold
                            && !ring_bkg.empty();
    output.ice_ring_score = AzimuthalIntegrationProfile::IceRingScore(
            have_ring_bkg ? ring_bkg : profile.GetResult1D(), integration.GetQBinCount(),
            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;
    // An unset limit masks nothing at that end. At the high-resolution end 0 does that on its own - no
    // pixel has d < 0, and the detector's own edge is where the pixels stop anyway; at the low-resolution
    // end every pixel lies above any finite stand-in, so it takes an infinite one.
    const float high_res = mask_high_res.value_or(0.0f);
    const float low_res = mask_low_res.value_or(INFINITY);
    auto const &resolution_map = integration.Resolution();
    for (int i = 0; i < mask_resolution.size(); i++)
        mask_resolution[i] = (resolution_map[i] > 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);
}