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

#include "MultiLatticeSearch.h"

#include <Eigen/Dense>

namespace {
    // Direct lattice vectors as matrix columns: M = [a | b | c]
    Eigen::Matrix3d LatticeMatrix(const CrystalLattice &latt) {
        const Coord a = latt.Vec0();
        const Coord b = latt.Vec1();
        const Coord c = latt.Vec2();

        Eigen::Matrix3d M;
        M.col(0) = Eigen::Vector3d(a.x, a.y, a.z);
        M.col(1) = Eigen::Vector3d(b.x, b.y, b.z);
        M.col(2) = Eigen::Vector3d(c.x, c.y, c.z);
        return M;
    }

    // Proper rotation mapping reference -> target (target = R * reference).
    // R = Mtarget * Mref^-1, projected to the nearest rotation via SVD.
    Eigen::Matrix3d RotationRefToTarget(const CrystalLattice &reference, const CrystalLattice &target) {
        Eigen::Matrix3d R = LatticeMatrix(target) * LatticeMatrix(reference).inverse();

        Eigen::JacobiSVD<Eigen::Matrix3d> svd(R, Eigen::ComputeFullU | Eigen::ComputeFullV);
        R = svd.matrixU() * svd.matrixV().transpose();
        if (R.determinant() < 0.0) {
            Eigen::Matrix3d U = svd.matrixU();
            U.col(2) *= -1.0;
            R = U * svd.matrixV().transpose();
        }
        return R;
    }
}

std::vector<MultiLatticeSearchResult> MultiLatticeSearch(const std::vector<CrystalLattice> &lattices,
                                                         float dist_tolerance,
                                                         float angle_tolerance_deg) {
    std::vector<MultiLatticeSearchResult> ret;
    if (lattices.empty())
        return ret;

    const CrystalLattice &reference = lattices[0];
    const UnitCell ref_cell = reference.GetUnitCell();

    ret.push_back({reference, reference, Coord(0, 0, 0)});

    for (size_t i = 1; i < lattices.size(); i++) {
        const CrystalLattice &latt = lattices[i];

        if (!latt.GetUnitCell().is_close(ref_cell, dist_tolerance, angle_tolerance_deg))
            continue;

        const Eigen::Matrix3d R = RotationRefToTarget(reference, latt);

        const Eigen::AngleAxisd aa(R);
        const Eigen::Vector3d rod = aa.angle() * aa.axis();
        const Coord rotation_vector(static_cast<float>(rod.x()),
                                    static_cast<float>(rod.y()),
                                    static_cast<float>(rod.z()));

        // output_lattice = R * reference, built straight from the Eigen matrix
        // so it is exactly a proper rotation of the reference (full double precision).
        const Eigen::Matrix3d out = R * LatticeMatrix(reference);
        ret.push_back({
            latt,
            CrystalLattice(Coord(out(0, 0), out(1, 0), out(2, 0)),
                           Coord(out(0, 1), out(1, 1), out(2, 1)),
                           Coord(out(0, 2), out(1, 2), out(2, 2))),
            rotation_vector
        });
    }

    return ret;
}

std::vector<MultiLatticeSearchResult> MultiLatticeSearchReduced(const std::vector<CrystalLattice> &lattices,
                                                                float dist_tolerance,
                                                                float angle_tolerance_deg) {
    if (lattices.empty())
        return {};

    // Reduce lattice 0 to its primitive Niggli setting; that becomes the reference basis.
    std::vector<CrystalLattice> reduced;
    reduced.reserve(lattices.size());

    // Reduce every candidate the same way so metric comparison and rotation
    // are all done in the reduced setting.
    for (const auto & lattice : lattices)
        reduced.push_back(lattice.NiggliReduce());

    auto ret = MultiLatticeSearch(reduced, dist_tolerance, angle_tolerance_deg);

    // Keep input_lattice pointing at the original (un-reduced) lattices for debugging.

    return ret;
}