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

#include "RingOptimizer.h"
#include "ceres/ceres.h"

struct RingResidual {
    RingResidual(double x, double y, double lambda,
                 double pixel_size,
                 double expected_q)
            : obs_x(x), obs_y(y),
              lambda(lambda),
              pixel_size(pixel_size),
              expected_len_recip_sq(expected_q * expected_q / (4.0 * M_PI * M_PI)) {}

    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 x_lab = (T(obs_x) - center_x[0]) * T(pixel_size); // convert to mm
        T y_lab = (T(obs_y) - center_y[0]) * T(pixel_size);
        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;
};

RingOptimizer::RingOptimizer(const DiffractionGeometry& geom) : reference(geom) {}

DiffractionGeometry RingOptimizer::Run(const std::vector<RingOptimizerInput> &input) {
    // 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)),
                nullptr,
                &center_x,
                &center_y,
                &distance,
                &rot1,
                &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;
    ceres::Solver::Summary summary;

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

    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;
}