// 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 "../common/JFJochMath.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());
    // 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);
}