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

#include "Combine3D.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <map>
#include <string>
#include <tuple>
#include <vector>

namespace {
    // A reflection of one event is the same raw (h,k,l) on a run of frames no more than this many
    // apart - tolerating a single sub-threshold frame in the middle of the rocking curve. The same
    // hkl only recurs after a large rotation (~180 deg), far beyond this, so events stay separate.
    constexpr float MAX_FRAME_GAP = 2.0f;

    // Locates one partial inside the input outcomes.
    struct PartialRef {
        int outcome;
        int reflection;
    };
}

std::vector<IntegrationOutcome> CombineRotationObservations(
    const std::vector<IntegrationOutcome>& partials,
    const DiffractionExperiment& experiment,
    Logger* logger,
    const std::string& dump_path) {

    const double min_total_partiality = experiment.GetScalingSettings().GetMinPartiality();
    const double capture_uncertainty_coeff = experiment.GetScalingSettings().GetCaptureUncertaintyCoeff();

    std::ofstream dump;
    if (!dump_path.empty()) {
        dump.open(dump_path);
        dump << "# h k l I sigma d n_frames captured_fraction\n";
    }

    // 1. Bucket every usable, already-scaled partial by raw (h,k,l).
    std::map<std::tuple<int, int, int>, std::vector<PartialRef>> by_hkl;
    size_t n_partials = 0;
    for (int o = 0; o < static_cast<int>(partials.size()); ++o) {
        const auto& reflections = partials[o].reflections;
        for (int j = 0; j < static_cast<int>(reflections.size()); ++j) {
            const auto& r = reflections[j];
            if (!std::isfinite(r.image_scale_corr) || r.image_scale_corr <= 0.0f)
                continue;
            if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f)
                continue;
            by_hkl[{r.h, r.k, r.l}].push_back({o, j});
            ++n_partials;
        }
    }

    // 2. Output mirrors the input images (same geom/latt) so per-image masking still works; each
    //    full is later placed on its peak frame's outcome.
    std::vector<IntegrationOutcome> out(partials.size());
    for (size_t o = 0; o < partials.size(); ++o) {
        out[o].geom = partials[o].geom;
        out[o].latt = partials[o].latt;
        out[o].mosaicity_deg = partials[o].mosaicity_deg;
        // Carry the first-pass per-image scaling metadata forward so it survives into the
        // _process.h5 / image.dat (the combine otherwise dropped it, leaving NaN per-image scale).
        out[o].image_scale_g = partials[o].image_scale_g;
        out[o].image_scale_b_factor_Ang2 = partials[o].image_scale_b_factor_Ang2;
        out[o].image_scale_cc = partials[o].image_scale_cc;
        out[o].image_scale_cc_n = partials[o].image_scale_cc_n;
        out[o].image_scale_wedge_deg = partials[o].image_scale_wedge_deg;
    }

    auto reflection_of = [&](const PartialRef& p) -> const Reflection& {
        return partials[p.outcome].reflections[p.reflection];
    };

    std::map<int, int> multiplicity_histogram;   // frames-per-event -> count
    size_t n_events = 0;

    // 3. Within each hkl, split into contiguous-frame rocking events and weight-sum each into a full.
    for (auto& [hkl, refs] : by_hkl) {
        std::sort(refs.begin(), refs.end(), [&](const PartialRef& a, const PartialRef& b) {
            return reflection_of(a).image_number < reflection_of(b).image_number;
        });

        size_t i = 0;
        while (i < refs.size()) {
            size_t k = i + 1;
            float last_frame = reflection_of(refs[i]).image_number;
            while (k < refs.size()) {
                const float frame = reflection_of(refs[k]).image_number;
                if (frame - last_frame > MAX_FRAME_GAP)
                    break;
                last_frame = frame;
                ++k;
            }

            // Pool the background across the rocking curve. Each partial subtracted its OWN
            // independently-estimated per-frame background, so a weak full accumulates one background-
            // estimation variance per frame. The true background is flat over the few frames of one
            // event, so replace each partial's background by the event mean and correct its intensity
            // by n_bkg*(bkg - pooled), where n_bkg ~ sigma^2/bkg is the effective background-pixel
            // count (single-frame events are a no-op). This trims the weak-full background noise.
            double pooled_bkg = 0.0;
            int n_pool = 0;
            for (size_t m = i; m < k; ++m) {
                const float b = reflection_of(refs[m]).bkg;
                if (std::isfinite(b)) { pooled_bkg += b; ++n_pool; }
            }
            pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;
            auto pooled_I = [&](const Reflection& r) {
                // Effective background-pixel count = (sigma^2 - I)/bkg: Poisson variance is signal +
                // background counts, so this is ~the true pixel count for BOTH weak and strong
                // reflections (sigma^2/bkg would over-count strong ones, whose variance is signal-led,
                // and over-correct them).
                const double n_bkg = std::max(0.0, static_cast<double>(r.sigma) * r.sigma - r.I) / std::max(r.bkg, 1.0f);
                return static_cast<double>(r.I) + n_bkg * (static_cast<double>(r.bkg) - pooled_bkg);
            };

            // First pass: a plain inverse-variance mean seeds F (and the event's peak frame /
            // resolution / total captured rocking fraction).
            double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
            float d = NAN;
            int peak_outcome = refs[i].outcome;
            float peak_frame = reflection_of(refs[i]).image_number;
            float peak_partiality = -1.0f;
            // All partials of one hkl share the resolution, hence the ice-ring status; carry it to the full.
            const bool on_ice = reflection_of(refs[i]).on_ice_ring;
            for (size_t m = i; m < k; ++m) {
                const auto& r = reflection_of(refs[m]);
                const double sigma_corr = static_cast<double>(r.sigma) * r.image_scale_corr;
                const double w = 1.0 / (sigma_corr * sigma_corr);
                sum_w += w;
                sum_wI += w * pooled_I(r) * r.image_scale_corr;
                sum_partiality += r.partiality;
                if (r.partiality > peak_partiality) {
                    peak_partiality = r.partiality;
                    peak_outcome = refs[m].outcome;
                    peak_frame = r.image_number;
                }
                if (!std::isfinite(d) && std::isfinite(r.d) && r.d > 0.0f)
                    d = r.d;
            }
            double F = sum_wI / sum_w;

            // De-biased Poisson reweighting (Kabsch profile fit): each frame's variance is its
            // background noise corr^2*(sigma^2 - I) - the part of the counting variance not from
            // signal - plus the *model* signal corr*F shared across the event. Using the model signal
            // rather than the down-fluctuating observed I stops weak partials being over-weighted,
            // the over-weighting that inflates the merge error model. Weights depend on F, so iterate.
            // (Per-reflection rocking-curve recentring was tried here and overfits - it scatters the
            // weak high-res reflections and lowers both CC1/2 and ISa - so it is deliberately omitted.)
            for (int iter = 0; iter < 3; ++iter) {
                sum_w = 0.0;
                sum_wI = 0.0;
                for (size_t m = i; m < k; ++m) {
                    const auto& r = reflection_of(refs[m]);
                    const double corr = r.image_scale_corr;
                    const double I_corr = pooled_I(r) * corr;
                    const double sigma_corr = static_cast<double>(r.sigma) * corr;
                    const double bkg_var = sigma_corr * sigma_corr - corr * I_corr;
                    double var = std::max(0.0, bkg_var) + corr * std::max(0.0, F);
                    if (!(var > 0.0))
                        var = sigma_corr * sigma_corr;
                    const double w = 1.0 / var;
                    sum_w += w;
                    sum_wI += w * I_corr;
                }
                F = sum_wI / sum_w;
            }

            const int n_frames = static_cast<int>(k - i);
            ++multiplicity_histogram[n_frames];
            ++n_events;
            i = k;

            // The total captured rocking fraction replaces the per-partial min_partiality cut: an
            // event seen at only a few percent of its curve is unreliable however many frames it spans.
            if (sum_w <= 0.0 || sum_partiality < min_total_partiality)
                continue;

            // Incomplete rocking-curve capture (sum_partiality = f < 1) leaves the full
            // extrapolated; charge the unobserved fraction (1-f) of the intensity as a systematic
            // uncertainty on top of the counting sigma, so the merge down-weights these
            // over-extrapolated (high-biased) low-capture fulls and the error model treats their
            // scatter as expected. coeff = 0 reproduces the plain counting sigma (baseline).
            double sigma_full = 1.0 / std::sqrt(sum_w);
            if (capture_uncertainty_coeff > 0.0) {
                const double f = std::min(1.0, sum_partiality);
                const double extra = capture_uncertainty_coeff * (1.0 - f) * std::max(0.0, F);
                sigma_full = std::sqrt(sigma_full * sigma_full + extra * extra);
            }

            Reflection full{};
            full.h = std::get<0>(hkl);
            full.k = std::get<1>(hkl);
            full.l = std::get<2>(hkl);
            full.I = static_cast<float>(F);
            full.sigma = static_cast<float>(sigma_full);
            full.d = d;
            full.image_number = peak_frame;
            full.partiality = 1.0f;        // a combined full, not a slice
            full.image_scale_corr = 1.0f;  // already fully corrected
            full.rlp = 1.0f;
            full.observed = true;
            full.on_ice_ring = on_ice;
            out[peak_outcome].reflections.push_back(full);

            if (dump.is_open())
                dump << full.h << ' ' << full.k << ' ' << full.l << ' ' << full.I << ' '
                     << full.sigma << ' ' << d << ' ' << n_frames << ' ' << sum_partiality << ' '
                     << static_cast<int>(peak_frame) << '\n';
        }
    }

    if (logger != nullptr) {
        size_t multi_frame = 0;
        double frame_sum = 0.0;
        for (const auto& [frames, count] : multiplicity_histogram) {
            frame_sum += static_cast<double>(frames) * count;
            if (frames >= 2)
                multi_frame += count;
        }
        const double mean_frames = n_events > 0 ? frame_sum / static_cast<double>(n_events) : 0.0;
        const double multi_pct = n_events > 0 ? 100.0 * multi_frame / static_cast<double>(n_events) : 0.0;
        logger->Info("3D combine: {} events from {} partials, mean {:.2f} frames/event, {:.1f}% multi-frame",
                     n_events, n_partials, mean_frames, multi_pct);
    }

    return out;
}