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

#include "SimpleRotXtalOptimizer.h"

#include <Eigen/Dense>

static bool SimpleRotXtalOptimizerInternal(
    SimpleRotXtalOptimizerData &data,
    const std::vector<SpotToSave> &spots,
    float tolerance,
    Eigen::Matrix3d &r_total)
{
    Coord vec0 = data.latt.Vec0();
    Coord vec1 = data.latt.Vec1();
    Coord vec2 = data.latt.Vec2();

    const double tol_sq = tolerance * tolerance;

    std::vector<Eigen::Vector3d> obs;
    std::vector<Eigen::Vector3d> exp;

    // Collect reflections
    for (const auto &pt : spots) {
        if (!data.index_ice_rings && pt.ice_ring)
            continue;

        Coord recip = data.geom.DetectorToRecip(pt.x, pt.y);

        double h_fp = recip * vec0;
        double k_fp = recip * vec1;
        double l_fp = recip * vec2;

        double h = std::round(h_fp);
        double k = std::round(k_fp);
        double l = std::round(l_fp);

        double d2 =
            (h - h_fp)*(h - h_fp) +
            (k - k_fp)*(k - k_fp) +
            (l - l_fp)*(l - l_fp);

        if (d2 > tol_sq)
            continue;

        obs.emplace_back(h_fp, k_fp, l_fp);
        exp.emplace_back(h, k, l);
    }

    if (obs.size() < data.min_spots)
        return false;

    const int N = static_cast<int>(obs.size());

    // Build linear system A * omega = b
    Eigen::MatrixXd A(3*N, 3);
    Eigen::VectorXd b(3*N);

    for (int i = 0; i < N; i++) {
        const auto &h = obs[i];
        const auto &e = exp[i];

        // Cross-product matrix of h_obs: [h]× such that [h]× · ω = h × ω
        // But we need ω × h = -h × ω, so use -[h]×
        A.row(3*i + 0) <<   0.0,  h.z(), -h.y();
        A.row(3*i + 1) << -h.z(),  0.0,   h.x();
        A.row(3*i + 2) <<  h.y(), -h.x(),  0.0;

        b.segment<3>(3*i) = e - h;
    }

    // Solve least squares
    Eigen::Vector3d omega = A.colPivHouseholderQr().solve(b);

    // Optional sanity check
    if (!omega.allFinite())
        return false;

    // Apply incremental rotation and accumulate it
    const double angle = omega.norm();
    Eigen::Matrix3d r_inc = Eigen::Matrix3d::Identity();
    if (angle > 0.0)
        r_inc = Eigen::AngleAxisd(angle, omega / angle).toRotationMatrix();

    r_total = r_inc * r_total;

    if (angle > 0.0) {
        // Build the current lattice matrix (columns = a, b, c)
        Eigen::Matrix3d L;
        L.col(0) = Eigen::Vector3d(vec0.x, vec0.y, vec0.z);
        L.col(1) = Eigen::Vector3d(vec1.x, vec1.y, vec1.z);
        L.col(2) = Eigen::Vector3d(vec2.x, vec2.y, vec2.z);

        // The HKL-space rotation R transforms: h_new = R * h_old
        // This corresponds to: L_new = L_old * R^T  (real-space lattice)
        Eigen::Matrix3d L_new = L * r_inc.transpose();

        data.latt = CrystalLattice(
            Coord(L_new(0,0), L_new(1,0), L_new(2,0)),
            Coord(L_new(0,1), L_new(1,1), L_new(2,1)),
            Coord(L_new(0,2), L_new(1,2), L_new(2,2))
        );
    }

    return true;
}

bool SimpleRotXtalOptimizer(SimpleRotXtalOptimizerData &data, const std::vector<SpotToSave> &spots) {
    Eigen::Matrix3d r_total = Eigen::Matrix3d::Identity();

    if (!SimpleRotXtalOptimizerInternal(data, spots, 0.3, r_total))
        return false;
    SimpleRotXtalOptimizerInternal(data, spots, 0.2, r_total);
    const bool ok = SimpleRotXtalOptimizerInternal(data, spots, 0.1, r_total);

    Eigen::AngleAxisd aa(r_total);
    data.rotation[0] = aa.axis().x();
    data.rotation[1] = aa.axis().y();
    data.rotation[2] = aa.axis().z();
    data.angle = aa.angle() * 180.0 / M_PI;
    return ok;
}