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

#include <cmath>
#include <fstream>

#include <spdlog/fmt/fmt.h>

#include "RugnuxCalibration.h"
#include "../common/GitInfo.h"
#include "../image_analysis/geom_refinement/AssignSpotsToRings.h"
#include "../image_analysis/geom_refinement/RingOptimizer.h"
#include "../image_analysis/geom_refinement/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) {
    CalibrationResult result;
    result.geometry = fitted;

    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

CalibrationResult CalibrateFromProfile(const std::vector<float> &profile,
                                       const AzimuthalIntegrationMapping &mapping,
                                       const DiffractionGeometry &geom,
                                       const std::vector<float> &calibrant_ring_q) {
    const auto points = RingsFromAzimuthalProfile(profile, mapping, geom, calibrant_ring_q);
    if (points.empty())
        throw JFJochException(JFJochExceptionCategory::CalibrationError,
                              "No powder ring found in the summed azimuthal profile");
    return Summarize(RingOptimizer(geom).Run(points), points);
}

CalibrationResult CalibrateFromSpots(const std::vector<SpotToSave> &spots,
                                     const DiffractionGeometry &geom,
                                     const std::vector<float> &calibrant_ring_q) {
    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.
    GuessGeometry(fitted, spots, calibrant_ring_q);
    OptimizeGeometry(fitted, spots, calibrant_ring_q);
    return Summarize(fitted, AssignSpotsToRings(fitted, spots, calibrant_ring_q));
}

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());
    f << "poni_version: 2\n";
    f << "Detector: Detector\n";
    f << fmt::format("Detector_config: {{\"pixel1\": {:g}, \"pixel2\": {:g}, \"max_shape\": [{}, {}]}}\n",
                     pixel_m, pixel_m, experiment.GetYPixelsNumConv(), experiment.GetXPixelsNumConv());
    f << fmt::format("Distance: {:.9g}\n", geom.GetDetectorDistance_mm() * 1e-3);
    f << fmt::format("Poni1: {:.9g}\n", geom.GetBeamY_pxl() * pixel_m + half_pixel_m);
    f << fmt::format("Poni2: {:.9g}\n", geom.GetBeamX_pxl() * pixel_m + half_pixel_m);
    // rot2 and rot3 change SIGN on the way out, and rot1 does not. pyFAI has the slow axis increasing
    // BOTTOM to TOP; we use the MX convention, top to bottom. The two frames therefore differ by a
    // reflection in y, and conjugating a rotation by a reflection gives R(n, theta) -> R(Mn, -theta).
    // For rot1 the axis IS y, so the axis reverses and the sense reverses and the two cancel; for rot2
    // (about x) and rot3 (about the beam) the axis lies in the mirror plane, so only the sense reverses.
    // The angles mean the same thing in both frames - it is only the handedness of the frame that
    // differs - and the same two flips would apply on the way IN if a PONI file were ever read.
    // Poni1/Poni2 need no such change: they are distances from the corner of the sensor along each
    // axis, which the direction the axis runs in does not affect.
    // Verified against pyFAI on a LaB6 image: written without the flip, the rings pyFAI integrates are
    // BROADER than with no tilt at all (peak 42 against 30, mean ring-position error 0.0045 1/A against
    // 0.0027); with it they sharpen to 132 and 0.0005.
    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", negate(geom.GetPoniRot2_rad()));
    f << fmt::format("Rot3: {:.9g}\n", negate(geom.GetPoniRot3_rad()));
    f << fmt::format("Wavelength: {:.9g}\n", geom.GetWavelength_A() * 1e-10);
    f.flush();
    if (!f)
        throw JFJochException(JFJochExceptionCategory::FileWriteError, "Error writing " + path);
}