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

#pragma once

// The reference the LM solver of XtalRefine is checked against: the same XtalRefineProblem handed to
// Ceres exactly as XtalOptimizer used to build it - the same residual functors, losses, priors, bounds,
// manifold and options.

#include "ceres/ceres.h"
#include "../image_analysis/geom_refinement/XtalRefine.h"

struct XtalRefineCeresPrior {
    XtalRefineCeresPrior(double gx, double gy, double p0, double weight)
        : gx(gx), gy(gy), p0(p0), weight(weight) {}
    template<typename T>
    bool operator()(const T *const p, T *residual) const {
        residual[0] = T(weight) * (T(gx) * p[0] + T(gy) * p[1] - T(p0));
        return true;
    }
    double gx, gy, p0, weight;
};

inline ceres::Solver::Summary SolveXtalRefineCeres(XtalRefineProblem &p, int num_threads) {
    std::vector<XtalFrameConstants> frame_const;
    frame_const.reserve(p.frame_angle_rad.size());
    if (p.beam_and_orientation_only)
        for (const double angle: p.frame_angle_rad)
            frame_const.emplace_back(p.detector_rot, p.rot_vec, angle, p.latt_vec1, p.latt_vec2, p.crystal_system);

    ceres::Problem problem;
    for (size_t i = 0; i < p.residuals.size(); i++) {
        ceres::LossFunction *loss = p.weight_sq.empty()
                                        ? nullptr
                                        : new ceres::ScaledLoss(nullptr, p.weight_sq[i], ceres::TAKE_OWNERSHIP);
        if (p.beam_and_orientation_only)
            problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<XtalResidualBeamOrientation, 3, 2, 3>(
                    new XtalResidualBeamOrientation(p.residuals[i], p.distance_mm, frame_const[p.frame[i]])),
                loss, p.beam, p.latt_vec0);
        else
            problem.AddResidualBlock(
                new ceres::AutoDiffCostFunction<XtalResidualFixedDistance, 3, 2, 2, 3, 3, 3, 3>(
                    new XtalResidualFixedDistance(p.residuals[i], p.distance_mm)),
                loss, p.beam, p.detector_rot, p.rot_vec, p.latt_vec0, p.latt_vec1, p.latt_vec2);
    }
    for (const auto &prior: p.priors)
        problem.AddResidualBlock(
            new ceres::AutoDiffCostFunction<XtalRefineCeresPrior, 1, 2>(
                new XtalRefineCeresPrior(prior.gx, prior.gy, prior.p0, prior.weight)),
            nullptr, prior.block == XtalRefinePrior::Block::Beam ? p.beam : p.detector_rot);

    const auto bounds = [&](double *block, const double *lo, const double *hi, int n) {
        for (int i = 0; i < n; i++) {
            if (lo[i] > -XtalRefineProblem::kNoBound)
                problem.SetParameterLowerBound(block, i, lo[i]);
            if (hi[i] < XtalRefineProblem::kNoBound)
                problem.SetParameterUpperBound(block, i, hi[i]);
        }
    };
    if (p.beam_constant)
        problem.SetParameterBlockConstant(p.beam);
    if (!p.beam_and_orientation_only) {
        if (p.detector_rot_constant)
            problem.SetParameterBlockConstant(p.detector_rot);
        else
            bounds(p.detector_rot, p.detector_rot_lower, p.detector_rot_upper, 2);
        if (p.rot_vec_constant)
            problem.SetParameterBlockConstant(p.rot_vec);
        else
            problem.SetManifold(p.rot_vec, new ceres::SphereManifold<3>);
        if (p.latt_vec1_constant)
            problem.SetParameterBlockConstant(p.latt_vec1);
        else
            bounds(p.latt_vec1, p.latt_vec1_lower, p.latt_vec1_upper, 3);
        if (p.latt_vec2_constant)
            problem.SetParameterBlockConstant(p.latt_vec2);
        else
            bounds(p.latt_vec2, p.latt_vec2_lower, p.latt_vec2_upper, 3);
    }

    ceres::Solver::Options options;
    options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY;
    options.minimizer_progress_to_stdout = false;
    options.max_num_iterations = p.options.max_iterations;
    options.max_solver_time_in_seconds = p.options.max_time_s;
    options.logging_type = ceres::LoggingType::SILENT;
    options.num_threads = num_threads;
    ceres::Solver::Summary summary;
    ceres::Solve(options, &problem, &summary);
    return summary;
}