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

#include <algorithm>
#include <cmath>
#include <fstream>

#include <spdlog/fmt/fmt.h>

#include "PowderCalibration.h"
#include "../../common/GitInfo.h"
#include "../../common/JFJochMath.h"
#include "AssignSpotsToRings.h"   // FindCircleCenter
#include "RingOptimizer.h"
#include "RingsFromProfile.h"

namespace {

// How well the ring points sit on the fitted rings, in a unit a user can judge: the radial distance in
// pixels between where a point is and where the fitted geometry puts its ring. The fit's own residual
// is in q, so it is divided by the local dq/dr - measured by stepping one pixel outward along the radius
// rather than assumed, since dq/dr varies with two-theta and with the tilt.
//
// The beam centre enters a ring's radius as r(phi) = R + dx cos(phi) + dy sin(phi), so fitting it to n
// points of scatter s leaves the textbook var = 2 s^2 / n on each of dx and dy. That is the number that
// separates a beam centre that was measured from one that was merely reported.
CalibrationResult Summarize(const DiffractionGeometry &fitted,
                            const std::vector<RingOptimizerInput> &points,
                            const RingFitUncertainty &unc) {
    CalibrationResult result;
    result.geometry = fitted;
    result.uncertainty = unc;

    const float cx = fitted.GetBeamX_pxl();
    const float cy = fitted.GetBeamY_pxl();
    double sum_sq = 0.0;
    for (const auto &p : points) {
        const float r = std::hypot(p.x - cx, p.y - cy);
        if (!(r > 1.0f))
            continue;
        const float q = fitted.PxlToQ(p.x, p.y);
        const float dq_dr = fitted.PxlToQ(p.x + (p.x - cx) / r, p.y + (p.y - cy) / r) - q;
        if (!(std::abs(dq_dr) > 0.0f))
            continue;
        const double dr = (q - p.q_expected) / dq_dr;
        sum_sq += dr * dr;
        ++result.ring_points;
    }
    if (result.ring_points > 0) {
        result.rms_radial_pxl = std::sqrt(sum_sq / static_cast<double>(result.ring_points));
        result.beam_sigma_pxl = result.rms_radial_pxl
                * std::sqrt(2.0 / static_cast<double>(result.ring_points));
    }
    return result;
}

}   // namespace

float TiltSignificance(const DiffractionGeometry &geom, const RingFitUncertainty &unc) {
    if (!unc.valid)
        return 0.0f;
    float significance = 0.0f;
    if (unc.sigma_rot1_rad > 0.0)
        significance = std::max(significance,
                                std::abs(geom.GetPoniRot1_rad()) / static_cast<float>(unc.sigma_rot1_rad));
    if (unc.sigma_rot2_rad > 0.0)
        significance = std::max(significance,
                                std::abs(geom.GetPoniRot2_rad()) / static_cast<float>(unc.sigma_rot2_rad));
    return significance;
}

CalibrationResult CalibrateFromProfile(const std::vector<float> &profile,
                                       const AzimuthalIntegrationMapping &mapping,
                                       const DiffractionGeometry &geom,
                                       const std::vector<float> &calibrant_ring_q,
                                       bool refine_tilt,
                                       const std::vector<SpotToSave> &spots) {
    // Where to start. The ring search below is local - each ring is looked for inside a window a few
    // pixels of radius wide - so a header distance more than a percent or two out puts every ring
    // outside its own window, and what the fit then converges on is noise. Ask the rings what the
    // distance is rather than believing the header (see PowderAutoSeed.h), and because a powder pattern
    // has genuine distance aliases, ask for several answers and fit them all.
    const auto observed = RingRadiiFromProfile(profile, mapping, geom);
    const auto [radius_min, radius_max] = ProfileRadiusRange_pxl(mapping, geom);
    const auto candidates = CandidateDistancesFromPowderRings(observed, calibrant_ring_q, geom,
                                                              radius_min, radius_max);

    // ...and ask them where the beam is, for the same reason. The extraction looks for each ring within
    // a window a few pixels of radius wide, so a header centre more than about ten pixels out puts the
    // ring outside that window over much of the turn - and the fit then reads a cos(phi) signal off
    // whatever sectors are left, which is how a 20 px error used to end 31 px wrong. The rings answer
    // this without a calibrant and without a distance: a powder ring is a conic centred on the beam, so
    // a wrong centre makes EVERY ring's radius oscillate once per turn by the same amount.
    // Three places a beam centre can come from, and they fail in different regimes - which is the whole
    // reason to carry all of them. The header is right on a well-configured instrument and is what a
    // calibration is run to check. The rings' own wobble is exact while the error stays under about half
    // a ring spacing, and stops meaning anything beyond that, where each ring reaches the azimuthally
    // averaged profile as two horns rather than one peak. The circle through the spots reads nothing
    // from the header whatsoever, so it holds where both of the others have given up - measured, it
    // returns the same answer from a header 400 px out in the centre and eight times out in distance.
    // None of them is trusted; each is fitted and the residual chooses.
    std::vector<DiffractionGeometry> centres = {geom};
    const auto add_centre = [&](float x, float y) {
        // Within a pixel of one already on the list it is not another hypothesis, it is the same one,
        // and fitting it again would only double the work for two answers nothing can tell apart.
        for (const auto &c : centres)
            if (std::hypot(x - c.GetBeamX_pxl(), y - c.GetBeamY_pxl()) < 1.0f)
                return;
        DiffractionGeometry candidate = geom;
        centres.push_back(candidate.BeamX_pxl(x).BeamY_pxl(y));
    };

    if (const auto offset = BeamCentreOffsetFromProfile(profile, mapping, geom, observed))
        add_centre(geom.GetBeamX_pxl() + offset->first, geom.GetBeamY_pxl() + offset->second);

    // What the spots make of it, as a COMPLETE geometry rather than only a centre. Taking the circle
    // centre alone is not enough: the distance candidates are read off the azimuthally averaged profile,
    // which is averaged about the header's centre, and once that is badly wrong the profile no longer
    // shows rings at all - a ring smeared over hundreds of pixels cannot be un-smeared by reading its
    // bins differently. The spots have no such problem, because a spot's position is a fact about the
    // image; GuessGeometry votes for the circle centre, clusters the radii into rings and takes the
    // distance from the innermost one, which is exactly the path that makes --calibration spots immune
    // to the header. Let it produce one whole starting geometry of its own.
    std::optional<DiffractionGeometry> from_spots;
    if (spots.size() >= 3) {
        try {
            DiffractionGeometry guess = geom;
            GuessGeometry(guess, spots, calibrant_ring_q, refine_tilt);
            from_spots = guess;
        } catch (const std::exception &) {
            // No circle, or no ring clusters - the spots simply have nothing to say here. The profile's
            // own hypotheses stand on their own, so there is nothing to report and nothing to stop.
        }
    }

    // One starting distance, fitted to convergence. A seed only has to land in the fit's basin, not on
    // the answer: it is measured from blended peaks in the azimuthally-averaged profile and is good to
    // about a per cent, which is close enough to converge from but far enough to sit every search window
    // a few pixels off its ring - and an off-centre window takes its background off the ring's own
    // flank, which costs both points and residual. So re-extract at the geometry the fit converged to
    // and fit again. Nothing is re-read from disk, so the whole loop is free.
    struct Attempt {
        DiffractionGeometry geometry;
        std::vector<RingOptimizerInput> points;
        RingFitUncertainty uncertainty;
        double rms_radial_pxl = 0.0;
    };
    const auto fit_from = [&](const DiffractionGeometry &start, bool seeded,
                              bool tilt) -> std::vector<Attempt> {
        DiffractionGeometry current = start;
        Attempt attempt;
        constexpr int MAX_PASSES = 3;
        for (int pass = 0; pass < MAX_PASSES; ++pass) {
            // The profile's own geometry is the search on the first pass of an unseeded attempt, and
            // the geometry converged to on every pass after - which is where the rings have been
            // measured to be, rather than where the header guessed.
            const DiffractionGeometry *search = (pass == 0 && !seeded) ? nullptr : &current;
            auto pass_points = RingsFromAzimuthalProfile(profile, mapping, geom, calibrant_ring_q,
                                                         0.06f, 3.0f, search);
            if (pass_points.empty())
                break;

            RingFitUncertainty pass_unc;
            const auto pass_fitted = RingOptimizer(current, tilt).Run(pass_points, &pass_unc);
            const float moved = std::abs(pass_fitted.GetDetectorDistance_mm()
                                         - current.GetDetectorDistance_mm());
            attempt.points = std::move(pass_points);
            attempt.uncertainty = pass_unc;
            attempt.geometry = pass_fitted;
            current = pass_fitted;
            // A tenth of a micron of distance moves the outermost ring by far less than a thousandth of
            // a pixel, so there is nothing left for another pass to find.
            if (moved < 1e-4f)
                break;
        }
        if (attempt.points.empty())
            return {};

        // Two ways to take the final measurement, and they genuinely disagree about which is better.
        //
        // Following the rings sector by sector is what lets a badly placed beam centre be recovered at
        // all - a centre wrong by (dx, dy) puts a ring at a different q in every sector, and one window
        // centred on one q finds it only where that oscillation happens to be small. But it costs
        // precision once they have been found, for a reason worth stating: a window that moves with phi
        // makes every systematic of the peak finder - where the background line is taken, how the
        // centroid sits in the window - vary with phi too, and phi is exactly the axis the beam centre
        // is read off. Measured, rms 0.415 -> 0.525 px on a good 110 mm fit, and 0.831 when the window
        // followed the fitted tilt as well.
        //
        // The alternative is a window that is the same in every sector, which the binned geometry with
        // only its DISTANCE replaced gives by construction - a distance error moves every sector's
        // window equally, where a centre or a tilt error does not. That is more precise where it works
        // and finds nothing where the centre is far out. So do both and let the same rule that ranks
        // everything else here decide.
        std::vector<Attempt> out;
        out.push_back(attempt);

        DiffractionGeometry measure_geom = geom;
        measure_geom.DetectorDistance_mm(current.GetDetectorDistance_mm());
        auto measured = RingsFromAzimuthalProfile(profile, mapping, geom, calibrant_ring_q,
                                                  0.06f, 3.0f, &measure_geom);
        if (!measured.empty()) {
            Attempt fixed_window;
            RingFitUncertainty measured_unc;
            fixed_window.geometry = RingOptimizer(current, tilt).Run(measured, &measured_unc);
            fixed_window.points = std::move(measured);
            fixed_window.uncertainty = measured_unc;
            out.push_back(std::move(fixed_window));
        }

        for (auto &a : out)
            a.rms_radial_pxl = Summarize(a.geometry, a.points, a.uncertainty).rms_radial_pxl;
        return out;
    };

    // Every combination of what the rings said and what the header said, fitted, with the residual left
    // to choose. The header is a hypothesis like any other here, neither trusted nor discarded - and so
    // is the seeded beam centre, which is NOT simply better than the header's.
    //
    // The centre seed reads the once-per-turn wobble of the ring radii, and a tilt puts a term of that
    // same shape there too - one that grows as the ring's radius squared, where a centre error does not.
    // Pooling the rings into one offset therefore absorbs part of the tilt into the centre, which is
    // worth several pixels on a genuinely tilted detector and made a good 110 mm fit worse when it was
    // simply believed. What it buys is capture range, and only where the header centre is far enough out
    // that the extraction would otherwise find the rings over a fraction of the turn. Offering it as an
    // alternative start costs one more fit each and needs no rule about when it applies.
    struct Provenance { float seed_mm; double match; };
    struct Start { DiffractionGeometry geometry; bool tracked; Provenance provenance; };

    std::vector<Start> starts;
    // The spots' geometry is a start in its own right, not one to cross with the profile's distances -
    // its centre and its distance are measured together and belong together.
    if (from_spots)
        starts.push_back({*from_spots, true, {from_spots->GetDetectorDistance_mm(), 0.0}});
    for (size_t centre = 0; centre < centres.size(); ++centre) {
        for (size_t i = 0; i <= candidates.size(); ++i) {
            const bool seeded_distance = i < candidates.size();
            DiffractionGeometry start = centres[centre];
            if (seeded_distance)
                start.DetectorDistance_mm(candidates[i].distance_mm);
            starts.push_back({start, seeded_distance || centre > 0,
                              {seeded_distance ? candidates[i].distance_mm : 0.0f,
                               seeded_distance ? candidates[i].score : 0.0}});
        }
    }

    std::vector<std::pair<Attempt, Provenance>> attempts;
    for (const auto &start : starts)
        for (auto &attempt : fit_from(start.geometry, start.tracked, refine_tilt))
            attempts.emplace_back(std::move(attempt), start.provenance);

    // Residual alone cannot rank these: a starting distance so wrong that only one ring point survives
    // leaves a residual of exactly zero, and would win every time. How many ring measurements an
    // attempt explains is evidence in its own right, and the first thing to compare - an attempt that
    // accounts for half as much of the pattern is not in the running whatever it does with what is
    // left. Among those that explain a comparable amount, the residual decides, and decides clearly,
    // because only the true distance makes every ring fit at once - on the aliased LaB6 case the two
    // candidates differ by 0.4 px against 5.2 px, which is not a close call.
    size_t most_points = 0;
    for (const auto &[attempt, provenance] : attempts)
        most_points = std::max(most_points, attempt.points.size());

    std::optional<Attempt> best;
    Provenance best_provenance{};
    for (auto &[attempt, provenance] : attempts) {
        if (attempt.points.size() * 2 < most_points)
            continue;
        if (!best || attempt.rms_radial_pxl < best->rms_radial_pxl) {
            best = std::move(attempt);
            best_provenance = provenance;
        }
    }

    if (!best)
        throw JFJochException(JFJochExceptionCategory::CalibrationError,
                              "No powder ring found in the summed azimuthal profile");

    // Did the tilt earn its place? A single ring cannot separate a tilt from a beam-centre shift at all
    // - both move a ring's radius as cos(phi), and only how that amplitude grows with the ring's radius
    // tells them apart - and two barely-sampled rings cannot either. The fit still returns a tilt in
    // that case, because nothing stopped it, and pays for it with the beam centre: at 500 mm on the LaB6
    // series it swings to the opposite sign of every shorter distance, drags the PONI 28 px, and buys a
    // residual of 0.960 px against 0.962 px pinned. The covariance says so plainly - 0.1 sigma, and a
    // PONI quoted to +-180 px - so ask it, and where the answer is no, fit again with the tilt held.
    // The rule is on the geometry that gets REPORTED, not on what the last fit happened to do: a tilt
    // may only survive if the final fit had it free AND it cleared the test. Both halves are needed.
    // RingOptimizer pins the tilt by itself when every point it is given lies on one ring, and on a
    // barely-sampled pattern a later pass can land in exactly that state - which used to freeze the tilt
    // at whatever the FIRST pass produced and hand it back with sigma zero, i.e. an unmeasured tilt
    // wearing the appearance of a fixed one. Refitting from the header's tilt is what makes the
    // reported geometry honest in both cases.
    const bool tilt_was_free = best->uncertainty.valid
            && (best->uncertainty.sigma_rot1_rad > 0.0 || best->uncertainty.sigma_rot2_rad > 0.0);
    const float significance = tilt_was_free ? TiltSignificance(best->geometry, best->uncertainty) : 0.0f;
    bool tilt_refined = refine_tilt && tilt_was_free && significance >= TILT_MIN_SIGNIFICANCE;
    if (refine_tilt && !tilt_refined) {
        // From the header's tilt, not from the one being declined. fit_from now takes a whole geometry,
        // so without this the "pinned" refit would pin rot1/rot2 at exactly the unvalidated values the
        // gate just rejected - which is the same fault the gate exists to catch, reintroduced one level
        // up.
        DiffractionGeometry pinned_start = best->geometry;
        pinned_start.PoniRot1_rad(geom.GetPoniRot1_rad()).PoniRot2_rad(geom.GetPoniRot2_rad());
        std::optional<Attempt> pinned;
        for (auto &a : fit_from(pinned_start, true, false))
            if (!pinned || a.points.size() > pinned->points.size()
                || (a.points.size() == pinned->points.size()
                    && a.rms_radial_pxl < pinned->rms_radial_pxl))
                pinned = std::move(a);
        if (pinned)
            best = std::move(*pinned);
    }

    auto result = Summarize(best->geometry, best->points, best->uncertainty);
    if (from_spots) {
        result.spots_available = true;
        result.spots_beam_x_pxl = from_spots->GetBeamX_pxl();
        result.spots_beam_y_pxl = from_spots->GetBeamY_pxl();
        result.spots_disagreement_pxl =
                std::hypot(from_spots->GetBeamX_pxl() - best->geometry.GetBeamX_pxl(),
                           from_spots->GetBeamY_pxl() - best->geometry.GetBeamY_pxl());
    }
    result.tilt_refined = tilt_refined;
    result.tilt_significance = significance;
    result.seed_distance_mm = best_provenance.seed_mm;
    result.header_distance_mm = geom.GetDetectorDistance_mm();
    return result;
}

CalibrationResult CalibrateFromSpots(const std::vector<SpotToSave> &spots,
                                     const DiffractionGeometry &geom,
                                     const std::vector<float> &calibrant_ring_q,
                                     bool refine_tilt) {
    DiffractionGeometry fitted = geom;
    // From scratch (Hough circle centre + ring clustering), then refined: the guess pins the centre to a
    // whole pixel and only sees the spots its clustering kept, so the refine re-matches every spot at
    // that geometry.
    RingFitUncertainty unc;
    GuessGeometry(fitted, spots, calibrant_ring_q, refine_tilt);
    OptimizeGeometry(fitted, spots, calibrant_ring_q, refine_tilt, &unc);

    // The same question the profile path asks, for the same reason and with the same answer if no: a
    // tilt may only be reported if this fit had it free and it stands clear of its own uncertainty.
    // Pinning it means putting rot1/rot2 back where the geometry came in and refitting the beam centre
    // and distance around that, not simply deleting an angle from the answer.
    const bool tilt_was_free = unc.valid
            && (unc.sigma_rot1_rad > 0.0 || unc.sigma_rot2_rad > 0.0);
    const float significance = tilt_was_free ? TiltSignificance(fitted, unc) : 0.0f;
    const bool tilt_refined = refine_tilt && tilt_was_free && significance >= TILT_MIN_SIGNIFICANCE;
    if (refine_tilt && !tilt_refined) {
        fitted.PoniRot1_rad(geom.GetPoniRot1_rad()).PoniRot2_rad(geom.GetPoniRot2_rad());
        OptimizeGeometry(fitted, spots, calibrant_ring_q, false, &unc);
    }

    auto result = Summarize(fitted, AssignSpotsToRings(fitted, spots, calibrant_ring_q), unc);
    result.tilt_refined = tilt_refined;
    result.tilt_significance = significance;
    result.header_distance_mm = geom.GetDetectorDistance_mm();
    return result;
}

void WritePoniFile(const std::string &path, const DiffractionExperiment &experiment,
                   const DiffractionGeometry &geom) {
    std::ofstream f(path);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::FileWriteError, "Cannot write " + path);

    const double pixel_m = geom.GetPixelSize_mm() * 1e-3;
    // pyFAI's axis convention is the trap: Poni1 (and pixel1) is the SLOW axis - rows, our y - and
    // Poni2 the FAST axis - columns, our x - both in metres from the detector origin. A transposed PONI
    // file is silently wrong, so the mapping is spelled out here rather than left to the reader.
    //
    // DiffractionGeometry's beam_x/beam_y IS the PONI: LabCoord rotates the vector measured FROM that
    // pixel, i.e. it is the point of normal incidence, so it maps straight across with no correction.
    // GetDirectBeam_pxl() is a different quantity - where the direct beam lands - and parts from the
    // PONI as soon as rot1/rot2 are non-zero.
    //
    // The half pixel is the origin convention (see docs/DETECTOR_GEOMETRY.md): our coordinates are
    // pixel-centred, so beam_x = 948 means the CENTRE of pixel 948, while pyFAI measures from the edge
    // of the sensor and puts the centre of pixel i at (i + 0.5) * pixel size. Without it the pattern
    // pyFAI integrates sits half a pixel off ours.
    const double half_pixel_m = 0.5 * pixel_m;
    f << fmt::format("# Calibration done by Jungfraujoch rugnux {}\n", jfjoch_version());
    // poni_version 2.1 is what pyFAI introduced "orientation" with (pyFAI 2024.01).
    f << "poni_version: 2.1\n";
    f << "Detector: Detector\n";
    // orientation 2 is pyFAI's "origin at the top left of the image when looking FROM the sample",
    // which is the MX convention Jungfraujoch assembles to. Without it pyFAI assumes its own default,
    // orientation 3 (bottom left), and quietly believes increasing row means physically upwards. The
    // radial integration is identical either way - a mirror preserves 2theta - but the azimuth comes
    // out with the opposite sense, which matters for anything that uses chi (cake or sector
    // integration, texture).
    f << fmt::format("Detector_config: {{\"pixel1\": {:g}, \"pixel2\": {:g}, \"max_shape\": [{}, {}], "
                     "\"orientation\": 2}}\n",
                     pixel_m, pixel_m, experiment.GetYPixelsNumConv(), experiment.GetXPixelsNumConv());
    f << fmt::format("Distance: {:.9g}\n", geom.GetDetectorDistance_mm() * 1e-3);
    // Poni1 is measured from pyFAI's own origin, so declaring orientation 2 re-anchors it to the top
    // edge: the same physical point is now (height - 1 - beam_y) rows down from there.
    f << fmt::format("Poni1: {:.9g}\n",
                     (experiment.GetYPixelsNumConv() - 1 - geom.GetBeamY_pxl()) * pixel_m + half_pixel_m);
    f << fmt::format("Poni2: {:.9g}\n", geom.GetBeamX_pxl() * pixel_m + half_pixel_m);
    // With orientation declared, rot2 and rot3 change sign and rot1 does not, and Rot3 carries a
    // further half turn:
    //     (Rot1, Rot2, Rot3) = (+rot1, +rot2, -rot3 + pi)
    // A row flip is an improper transformation, so it reverses the sense of rotations about x and
    // about the beam while leaving the one about the vertical alone. The half turn is the azimuthal
    // reference: pyFAI's in-plane axes are the negatives of ours, so without it every chi comes out
    // 180 degrees away. It is a rotation about the beam, so it leaves 2theta untouched - which is
    // why radial integration was right all along and only the azimuth was wrong.
    // Pinned empirically against pyFAI 2026.5.0 on a tilted detector (4/-6.5/13 deg, off-centre
    // beam), against the lab positions of the NXmx chain: 2theta to 3.6e-15 deg and chi to 2.8e-14
    // deg over the whole detector. The half turn is needed for the orientation-3 form written before
    // this too, so it is not an artefact of declaring the orientation.
    f << fmt::format("Rot1: {:.9g}\n", geom.GetPoniRot1_rad());
    // negate() rather than a bare minus so an unrefined angle prints as 0 and not -0.
    const auto negate = [](float v) { return v == 0.0f ? 0.0f : -v; };
    f << fmt::format("Rot2: {:.9g}\n", geom.GetPoniRot2_rad());
    f << fmt::format("Rot3: {:.9g}\n", negate(geom.GetPoniRot3_rad()) + PI);
    f << fmt::format("Wavelength: {:.9g}\n", geom.GetWavelength_A() * 1e-10);
    f.flush();
    if (!f)
        throw JFJochException(JFJochExceptionCategory::FileWriteError, "Error writing " + path);
}