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

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

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

            // 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;
            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 * static_cast<double>(r.I) * 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 = static_cast<double>(r.I) * 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;

            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>(1.0 / std::sqrt(sum_w));
            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;
            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;
}