// 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 "RingsFromProfile.h"   // SectorPeakQ
#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

std::vector<float> ProfileRingTrack(float q_cal, const DiffractionGeometry &truth,
                                    const DiffractionGeometry &binned, int32_t azim_bins) {
    std::vector<float> track(std::max<int32_t>(azim_bins, 1), NAN);
    if (azim_bins < 1 || !(q_cal > 0.0f))
        return track;
    const float d = static_cast<float>(2.0 * PI) / q_cal;
    if (!(d > truth.GetWavelength_A() / 2.0f))
        return track;   // past the Ewald limit - this ring is on no detector

    // Walk the ring in `truth` finely enough that every sector is hit several times, average what lands
    // in each. Averaging rather than taking one sample per sector because the map from an azimuth in
    // `truth` to a sector of `binned` is not uniform once the two centres differ.
    std::vector<double> sum(track.size(), 0.0);
    std::vector<int> count(track.size(), 0);
    const int samples = 8 * azim_bins;
    for (int i = 0; i < samples; ++i) {
        const float phi = static_cast<float>(2.0 * PI * i / samples);
        const auto [x, y] = truth.ResPhiToPxl(d, phi);
        if (!std::isfinite(x) || !std::isfinite(y))
            continue;
        const float q = binned.PxlToQ(x, y);
        if (!std::isfinite(q) || !(q > 0.0f))
            continue;
        float phi_binned = binned.Phi_rad(x, y);
        if (!std::isfinite(phi_binned))
            continue;
        auto bin = static_cast<int32_t>(phi_binned / static_cast<float>(2.0 * PI)
                                        * static_cast<float>(azim_bins));
        bin = std::clamp<int32_t>(bin, 0, azim_bins - 1);
        sum[bin] += q;
        ++count[bin];
    }
    for (size_t i = 0; i < track.size(); ++i)
        if (count[i] > 0)
            track[i] = static_cast<float>(sum[i] / count[i]);
    return track;
}

std::optional<std::pair<float, float>> BeamCentreOffsetFromProfile(
        const std::vector<float> &profile,
        const AzimuthalIntegrationMapping &mapping,
        const DiffractionGeometry &geom,
        const std::vector<ObservedRingRadius> &observed) {
    const int32_t q_bins = mapping.GetQBinCount();
    const int32_t azim_bins = mapping.GetAzimuthalBinCount();
    // Two rings at least: each is searched in a window reaching half way to its nearest neighbour, so
    // a single ring has no neighbour to bound it and nothing says which ring a peak inside a boundless
    // window belongs to.
    if (azim_bins < 4 || q_bins < 8 || observed.size() < 2
        || profile.size() != static_cast<size_t>(q_bins) * static_cast<size_t>(azim_bins))
        return std::nullopt;

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

    // Radius of every q bin, so a peak found in q can be stated as a radius - which is the quantity the
    // cos(phi) law below is written in, and the only one that does not depend on the assumed distance.
    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);

    // Normal equations for r - R_j = dx cos(phi) + dy sin(phi), pooled over rings. R_j, each ring's own
    // mean radius, is eliminated by centring each ring's measurements on their own mean - which is why
    // no d-spacing and no distance is needed to get the centre out.
    double sxx = 0.0, sxy = 0.0, syy = 0.0, sxr = 0.0, syr = 0.0;
    int used_rings = 0;

    for (const auto &ring : observed) {
        // Search each ring about the radius the profile actually put it at, in a window reaching half
        // way to its neighbours - so this is looking for a ring it has already found rather than for one
        // a standard predicts, and a badly wrong beam centre cannot push it out of its own window.
        float gap = std::numeric_limits<float>::max();
        for (const auto &other : observed)
            if (&other != &ring)
                gap = std::min(gap, std::abs(other.radius_pxl - ring.radius_pxl));
        const float window_pxl = 0.5f * gap;
        if (!(window_pxl > 2.0f))
            continue;

        std::vector<float> measured(azim_bins, NAN);
        for (int32_t phi_bin = 0; phi_bin < azim_bins; ++phi_bin) {
            int lo = -1, hi = -1;
            for (int32_t i = 0; i < q_bins; ++i) {
                if (!std::isfinite(radius[i]))
                    continue;
                if (std::abs(radius[i] - ring.radius_pxl) <= window_pxl) {
                    if (lo < 0) lo = i;
                    hi = i;
                }
            }
            if (lo < 0 || hi - lo < 6)
                continue;
            const float q_obs = SectorPeakQ(profile, q_bins, phi_bin, lo, hi,
                                            low_q, q_spacing, 3.0f);
            if (!std::isfinite(q_obs))
                continue;
            measured[phi_bin] = MeanRingRadius_pxl(geom, q_obs);
        }

        // Half the sectors, and spread round the turn: dx and dy are read off a cos and a sin, so a ring
        // seen only on one side of the pattern constrains one combination of them and leaves the other
        // free. Requiring the mean of cos and of sin over the sectors used to be small is what says the
        // coverage is even enough for the pair to separate.
        double mean_r = 0.0, mean_c = 0.0, mean_s = 0.0;
        int n = 0;
        const auto phi_of = [&](int32_t k) {
            return (static_cast<double>(k) + 0.5) * 2.0 * PI / static_cast<double>(azim_bins);
        };
        for (int32_t k = 0; k < azim_bins; ++k) {
            if (!std::isfinite(measured[k])) continue;
            mean_r += measured[k];
            mean_c += std::cos(phi_of(k));
            mean_s += std::sin(phi_of(k));
            ++n;
        }
        if (n * 2 < azim_bins)
            continue;
        mean_r /= n; mean_c /= n; mean_s /= n;
        if (std::hypot(mean_c, mean_s) > 0.25)
            continue;

        for (int32_t k = 0; k < azim_bins; ++k) {
            if (!std::isfinite(measured[k])) continue;
            const double c = std::cos(phi_of(k)), s = std::sin(phi_of(k));
            const double r = measured[k] - mean_r;
            sxx += c * c; sxy += c * s; syy += s * s;
            sxr += c * r; syr += s * r;
        }
        ++used_rings;
    }

    if (used_rings == 0)
        return std::nullopt;
    const double det = sxx * syy - sxy * sxy;
    if (!(std::abs(det) > 1e-9))
        return std::nullopt;
    const double dx = (sxr * syy - syr * sxy) / det;
    const double dy = (syr * sxx - sxr * sxy) / det;
    if (!std::isfinite(dx) || !std::isfinite(dy))
        return std::nullopt;
    return std::make_pair(static_cast<float>(dx), static_cast<float>(dy));
}

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<DistanceCandidate> 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<DistanceCandidate> out;
    for (const int step : peaks) {
        if (out.size() >= max_candidates)
            break;
        const bool distinct = std::none_of(out.begin(), out.end(), [&](const DistanceCandidate &c) {
            return std::abs(grid[step] - c.distance_mm) < 0.05f * c.distance_mm;
        });
        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, score[step]});
    }
    return out;
}