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

#include "XtalOptimizer.h"
#include "ceres/ceres.h"

struct XtalResidual {
    XtalResidual(double x, double y, double lambda,
                 double pixel_size, double distance_mm,
                 double rot1, double rot2,
                 bool tetragonal,
                 double exp_h, double exp_k, double exp_l)
            : obs_x(x), obs_y(y),
              lambda(lambda),
              pixel_size(pixel_size),
              distance(distance_mm),
              rot1(rot1),
              rot2(rot2),
              c1(cos(rot1)),
              c2(cos(rot2)),
              s1(sin(rot1)),
              s2(sin(rot2)),
              tetragonal(tetragonal),
              exp_h(exp_h),
              exp_k(exp_k),
              exp_l(exp_l) {}

    template<typename T>
    bool operator()(const T* const center_x, const T* const center_y,
                    const T* const latt_vec0,
                    const T* const latt_vec1,
                    const T* const latt_vec2,
                    T* residual) const {

        // x_lab in mm
        T x_lab = (T(obs_x) - center_x[0]) * T(pixel_size);
        T y_lab = (T(obs_y) - center_y[0]) * T(pixel_size);
        T z_lab = T(distance);

        // apply rotations
        T x = x_lab * T(c1) + z_lab * T(s1);
        T y = y_lab * T(c2) + (-x_lab * T(s1) + z_lab * T(c1)) * T(s2);
        T z = -y_lab * T(s2) + (-x_lab * T(s1) + z_lab * T(c1)) * T(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);

        // HKL match expectation
        T pred_h = R_x * T(latt_vec0[0]) + R_y * T(latt_vec0[1]) + R_z * T(latt_vec0[2]);
        T pred_k = R_x * T(latt_vec1[0]) + R_y * T(latt_vec1[1]) + R_z * T(latt_vec1[2]);
        T pred_l = R_x * T(latt_vec2[0]) + R_y * T(latt_vec2[1]) + R_z * T(latt_vec2[2]);

        residual[0] = T(exp_h) - T(pred_h);
        residual[1] = T(exp_k) - T(pred_k);
        residual[2] = T(exp_l) - T(pred_l);
        if (tetragonal) {
            // Angles are 90 deg
            residual[3] = latt_vec0[0] * latt_vec1[0] + latt_vec0[1] * latt_vec1[1] + latt_vec0[2] * latt_vec1[2];
            residual[4] = latt_vec0[0] * latt_vec2[0] + latt_vec0[1] * latt_vec2[1] + latt_vec0[2] * latt_vec2[2];
            residual[5] = latt_vec1[0] * latt_vec2[0] + latt_vec1[1] * latt_vec2[1] + latt_vec1[2] * latt_vec2[2];

            T len_a = latt_vec0[0] * latt_vec0[0] + latt_vec0[1] * latt_vec0[1] + latt_vec0[2] * latt_vec0[2];
            T len_b = latt_vec1[0] * latt_vec1[0] + latt_vec1[1] * latt_vec1[1] + latt_vec1[2] * latt_vec1[2];

            residual[6] = len_a - len_b;
        } else {
            residual[3] = T(0.0);
            residual[4] = T(0.0);
            residual[5] = T(0.0);
            residual[6] = T(0.0);
        }
        return true;
    }

    const double obs_x, obs_y;
    const double lambda;
    const double pixel_size;
    const double exp_h;
    const double exp_k;
    const double exp_l;
    const double distance;
    const double rot1, rot2;
    const double c1,c2,s1,s2;
    const bool tetragonal;
};

void XtalOptimizer(XtalOptimizerData &data,
                   const std::vector<SpotToSave> &spots,
                   bool force_tetragonal) {
    try {
        Coord vec0 = data.latt.Vec0();
        Coord vec1 = data.latt.Vec1();
        Coord vec2 = data.latt.Vec2();

        // Initial guess for the parameters
        double center_x = data.geom.GetBeamX_pxl();
        double center_y = data.geom.GetBeamY_pxl();
        double latt_vec0[3] = {vec0.x, vec0.y, vec0.z};
        double latt_vec1[3] = {vec1.x, vec1.y, vec1.z};
        double latt_vec2[3] = {vec2.x, vec2.y, vec2.z};

        ceres::Problem problem;

        // Add residuals for each point
        for (const auto &pt: spots) {
            Coord recip = data.geom.DetectorToRecip(pt.x, pt.y);
            double h = std::round(recip * vec0);
            double k = std::round(recip * vec1);
            double l = std::round(recip * vec2);

            problem.AddResidualBlock(
                    new ceres::AutoDiffCostFunction<XtalResidual, 7, 1, 1, 3, 3, 3>(
                            new XtalResidual(pt.x, pt.y,
                                             data.geom.GetWavelength_A(),
                                             data.geom.GetPixelSize_mm(),
                                             data.geom.GetDetectorDistance_mm(),
                                             data.geom.GetPoniRot1_rad(),
                                             data.geom.GetPoniRot2_rad(),
                                             force_tetragonal,
                                             h, k, l)),
                    nullptr,
                    &center_x,
                    &center_y,
                    latt_vec0,
                    latt_vec1,
                    latt_vec2
            );
        }

        // 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);

        data.geom.BeamX_pxl(center_x).BeamY_pxl(center_y);
        data.latt = CrystalLattice(Coord(latt_vec0[0], latt_vec0[1], latt_vec0[2]),
                                   Coord(latt_vec1[0], latt_vec1[1], latt_vec1[2]),
                                   Coord(latt_vec2[0], latt_vec2[1], latt_vec2[2]));
    } catch (...) {
        // For convergence problems, we are just not doing anything...
    }
}