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

#include <algorithm>
#include <cmath>

#include "PowderAutoSeed.h"
#include "../../common/JFJochMath.h"

namespace {

// Radius in pixels of the ring at q, averaged over four azimuths. The average is the point: a beam
// centre that is wrong by (dx, dy) moves the ring's apparent radius by dx cos(phi) + dy sin(phi), which
// four azimuths a quarter turn apart cancel exactly. So these radii survive a wrong beam centre as well
// as a wrong distance, which is what lets the distance be measured before the centre is known.
float MeanRingRadius_pxl(const DiffractionGeometry &geom, float q) {
    const float cx = geom.GetBeamX_pxl();
    const float cy = geom.GetBeamY_pxl();
    // ResPhiToPxl THROWS past the Ewald limit rather than returning NaN, and a q range wide enough to
    // reach it is a setting, not a fault - so decide here instead of letting it out of the seed.
    if (!(q > 0.0f) || !(static_cast<float>(2.0 * PI) / q > geom.GetWavelength_A() / 2.0f))
        return NAN;
    double sum = 0.0;
    int n = 0;
    for (int i = 0; i < 4; ++i) {
        const auto [x, y] = geom.ResPhiToPxl(static_cast<float>(2.0 * PI) / q,
                                             static_cast<float>(i * PI / 2.0));
        if (!std::isfinite(x) || !std::isfinite(y))
            continue;
        sum += std::hypot(x - cx, y - cy);
        ++n;
    }
    return n > 0 ? static_cast<float>(sum / n) : NAN;
}

// The predicted radius of a ring at q, for a detector at D. Rings past the Ewald limit (q too large for
// this wavelength) have no radius at all and are reported as NaN rather than silently folded back.
float PredictedRadius_pxl(float q, float distance_mm, float wavelength_A, float pixel_mm) {
    const float sin_theta = wavelength_A * q / static_cast<float>(4.0 * PI);
    if (!(sin_theta > 0.0f) || sin_theta >= 1.0f)
        return NAN;
    const float two_theta = 2.0f * std::asin(sin_theta);
    // Past 90 degrees the ring is on the back of the detector, which is not a case a flat detector has.
    if (two_theta >= static_cast<float>(PI / 2.0))
        return NAN;
    return distance_mm * std::tan(two_theta) / pixel_mm;
}

}   // namespace

float ProfileQForRing(float q_cal, float d_true_mm, float d_binned_mm,
                      float wavelength_A, float pixel_mm) {
    const float r = PredictedRadius_pxl(q_cal, d_true_mm, wavelength_A, pixel_mm);
    if (!std::isfinite(r) || !(d_binned_mm > 0.0f) || !(wavelength_A > 0.0f))
        return NAN;
    // Straight back through the flat-detector relation the binning used. The tilt is not carried: at
    // seeding time it is whatever the header says, which is zero for every header that has not already
    // been calibrated, and a tenth of a degree moves a ring by well under the search window.
    const float two_theta = std::atan(r * pixel_mm / d_binned_mm);
    return static_cast<float>(4.0 * PI) * std::sin(0.5f * two_theta) / wavelength_A;
}

std::pair<float, float> ProfileRadiusRange_pxl(const AzimuthalIntegrationMapping &mapping,
                                               const DiffractionGeometry &geom) {
    const auto &settings = mapping.Settings();
    const float lo = MeanRingRadius_pxl(geom, settings.GetLowQ_recipA());
    const float hi = MeanRingRadius_pxl(geom, settings.GetHighQ_recipA());
    if (!std::isfinite(lo) || !std::isfinite(hi))
        return {0.0f, 0.0f};
    return {std::min(lo, hi), std::max(lo, hi)};
}

std::vector<ObservedRingRadius> RingRadiiFromProfile(const std::vector<float> &profile,
                                                     const AzimuthalIntegrationMapping &mapping,
                                                     const DiffractionGeometry &geom,
                                                     float min_peak_over_noise) {
    std::vector<ObservedRingRadius> out;

    const int32_t q_bins = mapping.GetQBinCount();
    const int32_t azim_bins = mapping.GetAzimuthalBinCount();
    if (q_bins < 16 || azim_bins < 1
        || profile.size() != static_cast<size_t>(q_bins) * static_cast<size_t>(azim_bins))
        return out;

    // Average over azimuth. A ring is a ring at every azimuth, so this is the profile with the most
    // counts behind it; the sectors are only needed later, to tell the beam centre from the tilt.
    // Bins no pixel fell in are NaN and are left out of their own average rather than counted as zero,
    // which would dig a hole where a module gap crosses the ring.
    std::vector<float> radial(q_bins, NAN);
    for (int32_t i = 0; i < q_bins; ++i) {
        double sum = 0.0;
        int n = 0;
        for (int32_t j = 0; j < azim_bins; ++j) {
            const float v = profile[static_cast<size_t>(j) * q_bins + i];
            if (std::isfinite(v)) { sum += v; ++n; }
        }
        if (n > 0)
            radial[i] = static_cast<float>(sum / n);
    }

    const auto &settings = mapping.Settings();
    const float low_q = settings.GetLowQ_recipA();
    const float q_spacing = settings.GetQSpacing_recipA();

    // Peaks are found in RADIUS, not in q, and that is the whole reason this works without knowing the
    // distance. Bin i was filled by the pixels whose q under the binning geometry is q_i, which is to
    // say the pixels at radius MeanRingRadius(q_i) - so the radius of a bin is a fact about the
    // detector, identical whatever distance was assumed, while its q is not. A window measured in
    // pixels therefore means the same thing at every assumed distance; a window measured in bins does
    // not, and at a wrongly large distance the whole q axis compresses until neighbouring rings fall
    // inside one window and no peak is the largest in it.
    std::vector<float> radius(q_bins, NAN);
    for (int32_t i = 0; i < q_bins; ++i)
        radius[i] = MeanRingRadius_pxl(geom, low_q + (static_cast<float>(i) + 0.5f) * q_spacing);

    // Half the width of the window a peak has to dominate, in pixels of radius. A powder ring is a few
    // pixels wide; two rings closer than twice this are not separated, which is a real resolution limit
    // rather than a tuning knob.
    // ...but the window still has to hold enough BINS to have a background and a peak in it. How many
    // bins eight pixels spans depends on the assumed distance - the further away the detector is
    // assumed to be, the more the q axis compresses and the fewer bins cover the same piece of the
    // detector - so a window that is only physical would collapse below three bins a side at a wrongly
    // large distance and find nothing at all. Take whichever of the two is wider.
    constexpr float HALF_WIDTH_PXL = 8.0f;
    constexpr int HALF_WIDTH_MIN_BINS = 4;
    for (int32_t i = 0; i < q_bins; ++i) {
        if (!std::isfinite(radius[i]) || !std::isfinite(radial[i]))
            continue;
        int lo = i, hi = i;
        while (lo > 0 && std::isfinite(radius[lo - 1])
               && (radius[i] - radius[lo - 1] <= HALF_WIDTH_PXL || i - lo < HALF_WIDTH_MIN_BINS)) --lo;
        while (hi + 1 < q_bins && std::isfinite(radius[hi + 1])
               && (radius[hi + 1] - radius[i] <= HALF_WIDTH_PXL || hi - i < HALF_WIDTH_MIN_BINS)) ++hi;
        // Two background bins at each end and a peak between them is the least this can work with.
        if (hi - lo < 6 || !std::isfinite(radial[lo]) || !std::isfinite(radial[hi]))
            continue;

        const auto bkg_at = [&](int k) {
            const float t = static_cast<float>(k - lo) / static_cast<float>(hi - lo);
            return radial[lo] + t * (radial[hi] - radial[lo]);
        };

        bool is_max = true;
        for (int k = lo; k <= hi && is_max; ++k)
            if (std::isfinite(radial[k]) && radial[k] > radial[i]) is_max = false;
        if (!is_max)
            continue;

        const float height = radial[i] - bkg_at(i);
        if (!(height > 0.0f))
            continue;

        // The scatter of the window's own ends, as the noise this peak has to stand clear of - the same
        // measure SectorPeakQ uses, and for the same reason: an absolute cut would need a value per
        // detector and per exposure.
        float s = 0.0f;
        int n = 0;
        for (int k : {lo, lo + 1, hi - 1, hi}) {
            if (!std::isfinite(radial[k])) continue;
            const float r = radial[k] - bkg_at(k);
            s += r * r;
            ++n;
        }
        if (n == 0)
            continue;
        const float noise = std::sqrt(s / static_cast<float>(n));
        if (!(height > min_peak_over_noise * noise))
            continue;

        // Intensity-weighted centroid over the bins above half height, in radius - the same estimator
        // SectorPeakQ uses in q, and for the same reason: it needs no line shape.
        double sum_wr = 0.0, sum_w = 0.0;
        for (int k = i; k >= lo && radial[k] - bkg_at(k) >= 0.5f * height; --k) {
            const double w = radial[k] - bkg_at(k);
            sum_wr += w * radius[k];
            sum_w += w;
        }
        for (int k = i + 1; k <= hi && radial[k] - bkg_at(k) >= 0.5f * height; ++k) {
            const double w = radial[k] - bkg_at(k);
            sum_wr += w * radius[k];
            sum_w += w;
        }
        if (sum_w > 0.0)
            out.push_back({static_cast<float>(sum_wr / sum_w), height});
    }

    std::sort(out.begin(), out.end(),
              [](const ObservedRingRadius &a, const ObservedRingRadius &b) { return a.height > b.height; });
    return out;
}

std::vector<float> CandidateDistancesFromPowderRings(const std::vector<ObservedRingRadius> &observed,
                                                     const std::vector<float> &calibrant_ring_q,
                                                     const DiffractionGeometry &geom,
                                                     float radius_min_pxl, float radius_max_pxl,
                                                     size_t max_candidates) {
    if (observed.size() < 2 || calibrant_ring_q.empty() || !(radius_max_pxl > radius_min_pxl))
        return {};

    const float wavelength_A = geom.GetWavelength_A();
    const float pixel_mm = geom.GetPixelSize_mm();
    if (!(wavelength_A > 0.0f) || !(pixel_mm > 0.0f))
        return {};

    double total_weight = 0.0;
    for (const auto &o : observed)
        total_weight += o.height;
    if (!(total_weight > 0.0))
        return {};

    // Half a per cent of the radius, floored at two pixels: a ring's own width and the profile's bin
    // both scale with neither, so the looser of the two is what a match has to survive.
    const auto tolerance = [](float r) { return std::max(2.0f, 0.005f * r); };

    const auto score_at = [&](float distance) {
        std::vector<float> predicted;
        for (const float q : calibrant_ring_q) {
            const float r = PredictedRadius_pxl(q, distance, wavelength_A, pixel_mm);
            if (std::isfinite(r) && r >= radius_min_pxl && r <= radius_max_pxl)
                predicted.push_back(r);
        }
        if (predicted.empty())
            return 0.0;

        double explained = 0.0;
        for (const auto &o : observed) {
            float nearest = std::numeric_limits<float>::max();
            for (const float p : predicted)
                nearest = std::min(nearest, std::abs(p - o.radius_pxl));
            if (nearest < tolerance(o.radius_pxl))
                explained += o.height;
        }

        size_t seen = 0;
        for (const float p : predicted) {
            float nearest = std::numeric_limits<float>::max();
            for (const auto &o : observed)
                nearest = std::min(nearest, std::abs(p - o.radius_pxl));
            if (nearest < tolerance(p))
                ++seen;
        }

        // Both halves, multiplied. Only rewarding explained peaks would choose the shortest distance on
        // offer, where the predicted rings are packed so tightly that every peak has one within
        // tolerance; only rewarding seen rings would choose the longest, where a single predicted ring
        // sits on a single peak and nothing else is asked of it.
        return (explained / total_weight)
                * (static_cast<double>(seen) / static_cast<double>(predicted.size()));
    };

    // Scanned in log steps so the resolution is relative: 2000 steps over 20-2000 mm is 0.23% each,
    // which only has to be fine enough to land in a basin - the value is solved for below. A linear
    // scan would be needlessly fine at 2 m and too coarse at 30 mm.
    constexpr int STEPS = 2000;
    constexpr double MIN_MM = 20.0, MAX_MM = 2000.0;
    std::vector<double> score(STEPS + 1);
    std::vector<float> grid(STEPS + 1);
    for (int step = 0; step <= STEPS; ++step) {
        grid[step] = static_cast<float>(
                MIN_MM * std::pow(MAX_MM / MIN_MM, static_cast<double>(step) / STEPS));
        score[step] = score_at(grid[step]);
    }

    // Local maxima, best first. A basin is many steps wide, so taking the grid's maxima directly would
    // return the same distance three times over; candidates closer together than 5% are the same answer
    // and only the better one is kept.
    std::vector<int> peaks;
    for (int step = 1; step < STEPS; ++step)
        if (score[step] > 0.0 && score[step] >= score[step - 1] && score[step] > score[step + 1])
            peaks.push_back(step);
    std::sort(peaks.begin(), peaks.end(), [&](int a, int b) { return score[a] > score[b]; });

    std::vector<float> out;
    for (const int step : peaks) {
        if (out.size() >= max_candidates)
            break;
        const bool distinct = std::none_of(out.begin(), out.end(), [&](float d) {
            return std::abs(grid[step] - d) < 0.05f * d;
        });
        if (!distinct)
            continue;

        // The scan fixes the BASIN, not the value: its steps are 0.23% apart, which at 110 mm is a
        // quarter of a millimetre and enough to move the outer rings by more than a pixel. With the
        // pairing settled the distance is linear - r = D tan(2theta) / p with tan(2theta) known per
        // ring - so solve it outright over the pairs this basin matched, weighted by peak height.
        const float coarse = grid[step];
        double num = 0.0, den = 0.0;
        for (const auto &o : observed) {
            float nearest = std::numeric_limits<float>::max(), nearest_t = 0.0f;
            for (const float q : calibrant_ring_q) {
                const float r = PredictedRadius_pxl(q, coarse, wavelength_A, pixel_mm);
                if (!std::isfinite(r)) continue;
                if (std::abs(r - o.radius_pxl) < nearest) {
                    nearest = std::abs(r - o.radius_pxl);
                    nearest_t = r / coarse;   // the ring's radius per mm of distance
                }
            }
            if (nearest < tolerance(o.radius_pxl) && nearest_t > 0.0f) {
                num += static_cast<double>(o.height) * nearest_t * o.radius_pxl;
                den += static_cast<double>(o.height) * nearest_t * nearest_t;
            }
        }
        out.push_back(den > 0.0 ? static_cast<float>(num / den) : coarse);
    }
    return out;
}