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

#include <cstdlib>
#include "Rugnux.h"
#include "../image_analysis/structure_refinement/ModelValidation.h"
#include "../image_analysis/indexing/SpindleCuspLoss.h"
#include "../image_analysis/bragg_integration/SpotWidth.h"
#include "../image_analysis/bragg_integration/SpotFootprint.h"
#include "../image_analysis/bragg_integration/BraggStencil.h"   // BRAGG_FOOTPRINT_NSIGMA
#include "../image_analysis/hot_pixels/HotPixels.h"
#include "../image_analysis/SensorAbsorption.h"
#include "DiagnosticOutput.h"

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <condition_variable>
#include <deque>
#include <cstring>
#include <functional>
#include <future>
#include <mutex>
#include <numeric>
#include <map>
#include <set>
#include <sstream>
#include <thread>

#include "../reader/JFJochHDF5Reader.h"
#include "../common/JFJochMath.h"
#include "../common/ParallelFor.h"
#include "../common/Logger.h"
#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
#include "../common/CUDAWrapper.h"
#include "../common/JFJochException.h"
#include "../common/time_utc.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
#include "../image_analysis/beam_stop/ShadowFinder.h"
#include "../image_analysis/IndexAndRefine.h"
#include "../image_analysis/indexing/AnalyzeIndexing.h"
#include "../image_analysis/indexing/HarmonicContamination.h"
#include "../image_analysis/geom_refinement/BeamCenterFFT.h"
#include "../image_analysis/geom_refinement/BeamCenterFromBackground.h"
#include "../image_analysis/geom_refinement/GeometryRefiner.h"
#include "../image_analysis/indexing/IndexerThreadPool.h"
#include "../image_analysis/spot_finding/ImageSpotFinderCPU.h"
#include "../image_analysis/spot_finding/SpotUtils.h"
#include "gemmi/twin.hpp"
#include "../image_analysis/azint/AzIntEngineCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBuffer.h"
#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/image_preprocessing/ImagePreprocessorGPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"
#include "../image_analysis/spot_finding/AdaptiveSpotFinderGPU.h"
#endif
#include "../image_analysis/scale_merge/Merge.h"
#include "../image_analysis/scale_merge/RfreeFlags.h"
#include "../image_analysis/scale_merge/RotationScaleMerge.h"
#include "../image_analysis/scale_merge/ResolutionCutoff.h"
#include "../image_analysis/scale_merge/ReindexAmbiguity.h"
#include "../image_analysis/scale_merge/ScalingResult.h"
#include "../image_analysis/scale_merge/SearchSpaceGroup.h"
#include "../image_analysis/geom_refinement/PostRefine.h"
#include "../image_analysis/lattice_search/LatticeSearch.h"
#include "../image_analysis/lattice_search/LePageLattice.h"
#include "../image_analysis/scale_merge/AnisotropyAnalysis.h"
#include "../image_analysis/scale_merge/FrenchWilson.h"
#include "../image_analysis/scale_merge/TwinningAnalysis.h"
#include "../image_analysis/scale_merge/TranslationalNCS.h"
#include "../image_analysis/scale_merge/HKLKey.h"
#include "../image_analysis/scale_merge/ScaleOnTheFly.h"
#include "../image_analysis/scale_merge/StillsPartialityRefine.h"
#include "../image_analysis/WriteReflections.h"
#include "../image_analysis/bragg_integration/CalcISigma.h"
#include "../common/Definitions.h"
#include "../common/CorrelationCoefficient.h"
#include <array>
#include <map>
#include "RugnuxPipeline.h"

using namespace rugnux_internal;

std::optional<CalibrationResult> Rugnux::RebinAndRefit(const CalibrationResult &first,
                                                       int start_image, int end_image,
                                                       int images_to_process,
                                                       const AzimuthalIntegrationMapping &mapping,
                                                       const std::vector<SpotToSave> &spots) {
    Logger logger("Rugnux");
    const DiffractionGeometry before = experiment_.GetDiffractionGeometry();

    // How far a geometry moves a ring ON THE DETECTOR, against the radial width of one profile bin.
    // Below one bin, re-binning would put every pixel back where it already is. A beam-centre shift
    // moves a ring by its own size; a distance change moves it in proportion to the ring's radius.
    const auto [radius_min, radius_max] = ProfileRadiusRange_pxl(mapping, before);
    if (!(radius_max > radius_min))
        return std::nullopt;
    const float bin_pxl = (radius_max - radius_min) / static_cast<float>(mapping.GetQBinCount());
    const auto ring_gap_pxl = [&](const DiffractionGeometry &a, const DiffractionGeometry &b) {
        return std::hypot(a.GetBeamX_pxl() - b.GetBeamX_pxl(), a.GetBeamY_pxl() - b.GetBeamY_pxl())
                + radius_max * std::abs(a.GetDetectorDistance_mm() - b.GetDetectorDistance_mm())
                  / std::max(1e-3f, b.GetDetectorDistance_mm());
    };
    const auto ring_shift_pxl = [&](const DiffractionGeometry &g) { return ring_gap_pxl(g, before); };

    // Which geometries are worth re-binning about. The fitted one is the obvious candidate, but it was
    // itself fitted on the profile now believed mis-binned, so where the beam centre is badly out it is
    // wrong and re-binning about it only digs deeper - measured, rms 3.75 -> 5.66 on an exposure 200 px
    // off. The spots' geometry is the one that does not degrade there, having never read the header. So
    // try both where they differ and keep whichever comes back better; they cost a read each, and a
    // calibration run is a handful of images.
    std::vector<DiffractionGeometry> candidates;
    const auto add = [&](const DiffractionGeometry &g) {
        if (!(ring_shift_pxl(g) > bin_pxl))
            return;   // this geometry bins the images the way they are already binned
        for (const auto &c : candidates)
            if (ring_gap_pxl(g, c) < bin_pxl)
                return;   // and this one bins them the way another candidate already does
        candidates.push_back(g);
    };
    add(first.geometry);
    if (first.spots_geometry)
        add(*first.spots_geometry);

    if (candidates.empty()) {
        logger.Info("Powder calibration: the fitted geometry moves a ring by {:.2f} px, under the "
                    "{:.2f} px a profile bin spans - re-binning would change nothing",
                    ring_shift_pxl(first.geometry), bin_pxl);
        return std::nullopt;
    }

    std::optional<CalibrationResult> best;
    for (const auto &geometry : candidates) {
        if (cancelled_)
            break;
        logger.Info("Powder calibration: re-integrating binned about beam ({:.1f},{:.1f}) at {:.3f} mm, "
                    "which moves a ring {:.2f} px against a {:.2f} px bin",
                    geometry.GetBeamX_pxl(), geometry.GetBeamY_pxl(),
                    geometry.GetDetectorDistance_mm(), ring_shift_pxl(geometry), bin_pxl);

        // The mapping is the only thing the new geometry is needed for. Preprocessing does not depend on
        // where the beam is, so it keeps experiment_ and nothing here mutates the run's own geometry.
        DiffractionExperiment rebinned_experiment(experiment_);
        rebinned_experiment.BeamX_pxl(geometry.GetBeamX_pxl()).BeamY_pxl(geometry.GetBeamY_pxl())
                           .DetectorDistance_mm(geometry.GetDetectorDistance_mm())
                           .PoniRot1_rad(geometry.GetPoniRot1_rad())
                           .PoniRot2_rad(geometry.GetPoniRot2_rad());
        const AzimuthalIntegrationMapping rebinned_mapping(rebinned_experiment, pixel_mask_);

        // The sum over every image. The profile is a float sum, so the order it is added up in sets its
        // last bits: each fixed block of images is folded in image order into its own slot, and the
        // slots are added in block order, which gives the same profile at any thread count.
        std::vector<std::unique_ptr<AzimuthalIntegrationProfile>> block_sum(ReductionBlocks(images_to_process, 1));
        ParallelBlocks(images_to_process, std::max(1, config_.nthreads), [&](int b, int lo, int hi) {
            std::vector<uint8_t> decompression_buffer;
            JFJochReaderRawImage img;
            ImagePreprocessorCPU preprocessor(experiment_, pixel_mask_);
            ImagePreprocessorBuffer buffer(experiment_.GetPixelsNum());
            AzIntEngineCPU azint(rebinned_mapping);
            AzimuthalIntegrationProfile image_profile(rebinned_mapping);   // cleared by every azint.Run
            auto block = std::make_unique<AzimuthalIntegrationProfile>(rebinned_mapping);

            for (int ordinal = lo; ordinal < hi && !cancelled_; ++ordinal) {
                const int image_idx = start_image + ordinal * config_.stride;
                try {
                    if (!reader_.ReadRawImage(image_idx, img)) continue;
                    const uint8_t *image_ptr = img.image.GetUncompressedPtr(decompression_buffer);
                    preprocessor.Analyze(buffer, image_ptr, img.image.GetMode());
                    azint.Run(buffer, image_profile);
                    *block += image_profile;
                } catch (const std::exception &e) {
                    if (IsFatalResourceError(e)) throw;
                    logger.Error("Failed to re-integrate image {}: {}", image_idx, e.what());
                }
            }
            block_sum[b] = std::move(block);
        }, 1);
        AzimuthalIntegrationProfile summed(rebinned_mapping);
        for (const auto &block : block_sum)
            summed += *block;
        if (cancelled_)
            return std::nullopt;

        // The profile is measured about this geometry, so this is also what the fit starts from and what
        // its ring points are recovered against.
        try {
            auto refit = CalibrateFromProfile(summed.GetResult(), rebinned_mapping, geometry,
                                              config_.calibrant_ring_q,
                                              config_.calibration_refine_tilt, spots);
            logger.Info("Powder calibration: that gives {} ring points at rms {:.3f} px",
                        refit.ring_points, refit.rms_radial_pxl);
            // A refit that is not a measurement never beats one that is, whatever its residual. The
            // residual of a fit with a parameter pinned at a value nobody measured is not comparable
            // with the residual of one where every parameter was fitted - and this pass hands the
            // NEXT pass the geometry it pins to, so taking a non-measurement here launders it into
            // the following one's input.
            if (!best || (refit.converged && !best->converged)
                || (refit.converged == best->converged
                    && refit.ring_points * 2 >= best->ring_points
                    && refit.rms_radial_pxl < best->rms_radial_pxl))
                best = std::move(refit);
        } catch (const std::exception &e) {
            logger.Warning("Powder calibration: that profile could not be fitted ({})", e.what());
        }
    }

    if (!best)
        return std::nullopt;

    // Only if it is actually better. Re-binning about a geometry that was itself wrong makes the fit
    // worse, not better, and the first pass is then the answer to keep - it is measured on a profile
    // binned about a geometry now believed wrong, which is weak, but it is the one the run has.
    // Symmetric with the best-of fold above: a fit that is not a measurement never wins on residual,
    // and never loses on one either. The one-directional form threw away a CONVERGED re-fit whose rms
    // happened not to beat a non-converged first pass - which is the motivating case, where the first
    // pass declined the tilt and pinned the header's.
    if ((first.converged && !best->converged)
        || (first.converged == best->converged
            && (best->ring_points * 2 < first.ring_points
                || best->rms_radial_pxl >= first.rms_radial_pxl))) {
        logger.Info("Powder calibration: the re-binned fit is no better than the first "
                    "({} points at rms {:.3f} against {} at {:.3f}) - keeping the first pass",
                    best->ring_points, best->rms_radial_pxl, first.ring_points, first.rms_radial_pxl);
        return std::nullopt;
    }
    // The header the answer is compared against is the run's own, not the pass before it: a reader wants
    // to know how far the calibration moved the FILE.
    best->header_distance_mm = first.header_distance_mm;
    return best;
}

void Rugnux::FitCalibration(PipelineLocals &p, const std::vector<SpotToSave> &calibration_spot_list) {
    auto &logger = p.logger;
    auto &result = p.result;
    auto &observer = p.observer;
    auto &calibration_spots = p.calibration_spots;
    auto &start_image = p.start_image;
    auto &end_image = p.end_image;
    auto &images_to_process = p.images_to_process;
    auto &mapping = p.mapping;
    auto &plots = p.plots;
    if (observer)
        observer->OnPhase("Powder calibration");
    Step("Powder calibration");
    if (calibration_spots) {
        logger.Info("Powder calibration from {} pooled spots", calibration_spot_list.size());
        result.calibration = CalibrateFromSpots(calibration_spot_list,
                                                experiment_.GetDiffractionGeometry(),
                                                config_.calibrant_ring_q,
                                                config_.calibration_refine_tilt);
    } else {
        const auto *summed = plots.GetAzIntProfileObject();
        if (!summed)
            throw JFJochException(JFJochExceptionCategory::CalibrationError,
                                  "No azimuthal profile was accumulated");
        result.calibration = CalibrateFromProfile(summed->GetResult(), mapping,
                                                  experiment_.GetDiffractionGeometry(),
                                                  config_.calibrant_ring_q,
                                                  config_.calibration_refine_tilt,
                                                  calibration_spot_list);

        // Re-bin, and fit again on a profile that was binned about the answer.
        //
        // Everything above reads a profile built with the geometry the run STARTED from, and a
        // calibration is run precisely because that geometry is in doubt. Binning is not something
        // the fit can undo: which pixels landed in which (q, phi) bin was decided when the images
        // were read. Get the centre wrong and every ring is smeared across its own sectors - past a
        // hundred pixels the fit leaves rms 4.8 px even when it is started from the exact answer.
        // Get the distance wrong and the radial sampling is simply the wrong scale, which is why a
        // run recovering 110 mm from a 250 mm header ended at 1.1 px where the same data from a
        // right header gives 0.4.
        //
        // So integrate the images once more, binned with what was just fitted. A calibration run is
        // a handful of images and the fit is worth far more than the second read.
        const auto rebinned = RebinAndRefit(*result.calibration, start_image, end_image,
                                            images_to_process, mapping, calibration_spot_list);
        if (rebinned)
            result.calibration = rebinned;
    }
}