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

#pragma once

#include <limits>
#include <vector>

#include "XtalResidual.h"
#include "LMSolver.h"

// A soft restraint w * (g . p - p0) on one direction g of a two-component block - the beam centre or
// the detector tilt, which are the same gauge seen twice and so take the same direction (see
// XtalOptimizer).
struct XtalRefinePrior {
    enum class Block { Beam, DetectorRot } block = Block::Beam;
    double gx = 0.0, gy = 0.0, p0 = 0.0, weight = 0.0;
};

// The least-squares problem XtalOptimizer solves, as data: the residuals with the frame each belongs
// to, the parameter blocks with what is held and what is bounded, and the priors. The parameter arrays
// are in/out. Parameter blocks: beam(2), detector_rot(2), rot_vec(3, the spindle, refined on the sphere),
// latt_vec0(3, orientation angle-axis), latt_vec1(3, cell lengths), latt_vec2(3, cell angles) - see
// XtalResidual. The distance is a constant of the problem.
struct XtalRefineProblem {
    static constexpr double kNoBound = std::numeric_limits<double>::max();

    gemmi::CrystalSystem crystal_system = gemmi::CrystalSystem::Triclinic;
    // Only beam and orientation free, everything else held: the reduced residual (see
    // XtalResidualBeamOrientation), whose crystal half is a constant of the problem.
    bool beam_and_orientation_only = false;
    double distance_mm = 0.0;

    std::vector<XtalResidual> residuals;
    std::vector<int> frame;              // per residual, an index into frame_angle_rad
    std::vector<double> frame_angle_rad;
    std::vector<double> weight_sq;       // per residual; empty = unweighted

    double beam[2] = {0, 0};
    double detector_rot[2] = {0, 0};
    double rot_vec[3] = {1, 0, 0};
    double latt_vec0[3] = {0, 0, 0};
    double latt_vec1[3] = {0, 0, 0};
    double latt_vec2[3] = {0, 0, 0};

    bool beam_constant = false;
    bool detector_rot_constant = true;
    bool rot_vec_constant = true;
    bool latt_vec1_constant = true;
    bool latt_vec2_constant = true;

    double detector_rot_lower[2] = {-kNoBound, -kNoBound}, detector_rot_upper[2] = {kNoBound, kNoBound};
    double latt_vec1_lower[3] = {-kNoBound, -kNoBound, -kNoBound}, latt_vec1_upper[3] = {kNoBound, kNoBound, kNoBound};
    double latt_vec2_lower[3] = {-kNoBound, -kNoBound, -kNoBound}, latt_vec2_upper[3] = {kNoBound, kNoBound, kNoBound};

    std::vector<XtalRefinePrior> priors;

    LMOptions options;
};

// Solves the problem in place. The residual sums are cut into blocks that depend on the problem
// alone, so the answer is the same at any thread count.
LMSummary SolveXtalRefine(XtalRefineProblem &problem, int num_threads);