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

#include <cmath>
#include <algorithm>

#include "../../common/DetectorOrientation.h"
#include "../../common/JFJochMath.h"
#include "RingOptimizer.h"
#include "ceres/ceres.h"

struct RingResidual {
    RingResidual(double x, double y, double lambda,
                 double pixel_size,
                 double expected_q,
                 const DetectorOrientation &orientation)
            : obs_x(x), obs_y(y),
              lambda(lambda),
              pixel_size(pixel_size),
              expected_len_recip_sq(expected_q * expected_q / (4.0 * PI * PI)) {
        const RotMatrix delta = orientation.Matrix();
        det_m00 = delta.Column(0).x;
        det_m01 = delta.Column(1).x;
        det_m10 = delta.Column(0).y;
        det_m11 = delta.Column(1).y;
    }

    template<typename T>
    bool operator()(const T* const center_x, const T* const center_y,
                    const T* const distance, const T* const rot1,
                    const T* const rot2, T* residual) const {

        // Calculate lab coordinates from observed pixel coordinates
        T u_lab = (T(obs_x) - center_x[0]) * T(pixel_size); // convert to mm
        T v_lab = (T(obs_y) - center_y[0]) * T(pixel_size);
        // The discrete image orientation, which turns the offset from the PONI before the tilt acts.
        // Identity unless a detector says otherwise. It cannot change a ring's radius, but it does
        // change which way the tilt tips the ring, which is exactly what this fits.
        T x_lab = det_m00 * u_lab + det_m01 * v_lab;
        T y_lab = det_m10 * u_lab + det_m11 * v_lab;
        T z_lab = distance[0];

// Apply rotations around y and x axes
        T c1 = ceres::cos(rot1[0]);
        T c2 = ceres::cos(rot2[0]);
        T s1 = ceres::sin(rot1[0]);
        T s2 = ceres::sin(rot2[0]);

        T x = x_lab * c1 + z_lab * s1;
        T y = y_lab * c2 + (-x_lab * s1 + z_lab * c1) * s2;
        T z = -y_lab * s2 + (-x_lab * s1 + z_lab * c1) * c2;

        // convert to recip space
        T lab_norm = ceres::sqrt(x*x + y*y + z*z);

        T R_x = x / (lab_norm * T(lambda));
        T R_y = y / (lab_norm * T(lambda));
        T R_z = (z / lab_norm - T(1.0)) / T(lambda);

        T predicted_len_recip_sq = R_x * R_x + R_y * R_y + R_z * R_z;

        residual[0] = predicted_len_recip_sq - T(expected_len_recip_sq);
        return true;
    }

    const double obs_x, obs_y;
    const double lambda;
    const double pixel_size;
    const double expected_len_recip_sq;
    double det_m00, det_m01, det_m10, det_m11;
};

RingOptimizer::RingOptimizer(const DiffractionGeometry& geom, bool refine_tilt)
        : reference(geom), refine_tilt(refine_tilt) {}

namespace {

// Covariance of the converged fit, scaled to the residual scatter this fit actually left. Ceres hands
// back the bare (J^T J)^-1 of an unweighted problem, which carries the shape of the correlations but
// not their size; multiplying by chi2 per degree of freedom is what turns it into a sigma. The
// residual is in q^2, so the scale is in q^2 too - and it cancels, because (J^T J)^-1 is in
// parameter^2 per residual^2. The sigmas therefore come out in pixels, mm and radians whatever the
// residual is measured in, which is why this does not need to wait on the residual being restated.
void ComputeUncertainty(ceres::Problem &problem, const ceres::Solver::Summary &summary,
                        const std::vector<const double *> &free_blocks,
                        double *const centre_x, double *const centre_y, double *const distance,
                        double *const rot1, double *const rot2,
                        RingFitUncertainty &unc) {
    const int p = static_cast<int>(free_blocks.size());
    const int n = static_cast<int>(summary.num_residuals);
    unc.free_parameters = p;
    if (n <= p)
        return;

    // DENSE_SVD rather than the SPARSE_QR default: five parameters is not a sparse problem, and SVD is
    // the one that survives a near-degenerate normal matrix long enough to report how degenerate it is.
    ceres::Covariance::Options options;
    options.algorithm_type = ceres::DENSE_SVD;
    options.num_threads = 1;
    ceres::Covariance covariance(options);

    std::vector<std::pair<const double *, const double *>> pairs;
    for (const double *a : free_blocks)
        for (const double *b : free_blocks)
            pairs.emplace_back(a, b);
    // A failure here is not an error to report upwards: it means the problem is rank-deficient at the
    // solution, i.e. some direction in parameter space costs the fit nothing at all. valid = false
    // carries that, and it is the strongest statement this struct can make.
    if (!covariance.Compute(pairs, &problem))
        return;

    std::vector<double> cov(static_cast<size_t>(p) * p);
    if (!covariance.GetCovarianceMatrix(free_blocks, cov.data()))
        return;

    // 2 * final_cost is sum of squared residuals - Ceres' cost is half of it.
    unc.chi2_per_dof = 2.0 * summary.final_cost / static_cast<double>(n - p);
    for (auto &v : cov)
        v *= unc.chi2_per_dof;

    const auto index_of = [&](const double *block) {
        for (int i = 0; i < p; ++i)
            if (free_blocks[i] == block) return i;
        return -1;
    };
    const auto sigma = [&](const double *block) {
        const int i = index_of(block);
        return i < 0 ? 0.0 : std::sqrt(std::max(0.0, cov[i * p + i]));
    };
    const auto corr = [&](const double *a, const double *b) {
        const int i = index_of(a), j = index_of(b);
        if (i < 0 || j < 0) return 0.0;
        const double d = std::sqrt(cov[i * p + i] * cov[j * p + j]);
        return d > 0.0 ? cov[i * p + j] / d : 0.0;
    };

    unc.sigma_beam_x_pxl = sigma(centre_x);
    unc.sigma_beam_y_pxl = sigma(centre_y);
    unc.sigma_distance_mm = sigma(distance);
    unc.sigma_rot1_rad = sigma(rot1);
    unc.sigma_rot2_rad = sigma(rot2);
    unc.corr_beam_x_rot1 = corr(centre_x, rot1);
    unc.corr_beam_y_rot2 = corr(centre_y, rot2);
    unc.valid = true;
}

}   // namespace

DiffractionGeometry RingOptimizer::Run(const std::vector<RingOptimizerInput> &input,
                                       RingFitUncertainty *unc) {
    // Initial guess for the parameters
    double center_x = reference.GetBeamX_pxl();
    double center_y = reference.GetBeamY_pxl();
    double distance = reference.GetDetectorDistance_mm();
    double rot1 = reference.GetPoniRot1_rad();
    double rot2 = reference.GetPoniRot2_rad();

    ceres::Problem problem;

    // Add residuals for each point
    for (const auto& pt : input) {
        problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<RingResidual, 1, 1, 1, 1, 1, 1>(
                        new RingResidual(pt.x, pt.y,
                                         reference.GetWavelength_A(),
                                         reference.GetPixelSize_mm(),
                                         pt.q_expected,
                                         reference.GetOrientation())),
                nullptr,
                &center_x,
                &center_y,
                &distance,
                &rot1,
                &rot2
        );
    }

    // A single ring cannot tell the beam centre from the detector tilt: both displace its radius as
    // cos(phi), and what separates them is only how that amplitude scales with the ring's radius, which
    // takes two rings. Hold the tilt where it was given, so the one thing a single ring does fix - where
    // its centre lies - comes out rather than being traded away against an unconstrained tilt.
    const bool single_ring = !input.empty()
            && std::all_of(input.begin(), input.end(), [&](const RingOptimizerInput &p) {
                   return p.q_expected == input.front().q_expected;
               });
    const bool tilt_free = !(single_ring || !refine_tilt);
    if (!tilt_free) {
        problem.SetParameterBlockConstant(&rot1);
        problem.SetParameterBlockConstant(&rot2);
    }

    // Configure solver
    ceres::Solver::Options options;
    options.linear_solver_type = ceres::DENSE_QR;
    options.minimizer_progress_to_stdout = false;
    options.logging_type = ceres::LoggingType::SILENT;
    options.num_threads = 1;

    ceres::Solver::Summary summary;

    // Run optimization
    ceres::Solve(options, &problem, &summary);

    if (unc && summary.IsSolutionUsable()) {
        std::vector<const double *> free_blocks = {&center_x, &center_y, &distance};
        if (tilt_free) {
            free_blocks.push_back(&rot1);
            free_blocks.push_back(&rot2);
        }
        ComputeUncertainty(problem, summary, free_blocks,
                           &center_x, &center_y, &distance, &rot1, &rot2, *unc);
    }

    // A failed fit must not move the detector geometry. Both callers assign the result straight
    // back over the geometry they passed in, so handing back the reference leaves the calibration
    // where it was instead of committing a diverged beam centre and distance.
    if (!summary.IsSolutionUsable())
        return reference;

    DiffractionGeometry refined_geom(reference);
    refined_geom.BeamX_pxl(center_x).BeamY_pxl(center_y).DetectorDistance_mm(distance)
                .PoniRot1_rad(rot1).PoniRot2_rad(rot2);

    return refined_geom;
}