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

#include "MultiLatticeSearch.h"

#include <array>
#include <vector>

#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;
    }

    struct BasisOp {
        std::array<int, 3> perm;
        std::array<int, 3> signs; // each +1 or -1
    };

    // The 24 right-handed (det=+1) sign+permutation reparametrizations of a basis (a,b,c).
    // Identity (perm 0,1,2; signs +,+,+) is first so it's preferred when several ops match
    // the reference cell. For low-symmetry cells, only a handful pass is_close; for cubic,
    // all 24 may pass and the first-match rule keeps behavior deterministic.
    const std::vector<BasisOp> &RightHandedOps() {
        static const std::vector<BasisOp> ops = []() {
            std::vector<BasisOp> v;
            const std::array<std::array<int, 3>, 6> perms = {{
                {0, 1, 2}, {1, 2, 0}, {2, 0, 1}, // even
                {0, 2, 1}, {2, 1, 0}, {1, 0, 2}  // odd
            }};
            for (size_t p = 0; p < perms.size(); p++) {
                const int parity = (p < 3) ? +1 : -1;
                for (int s0 : {+1, -1})
                    for (int s1 : {+1, -1})
                        for (int s2 : {+1, -1})
                            if (parity * s0 * s1 * s2 == +1)
                                v.push_back({perms[p], {s0, s1, s2}});
            }
            return v;
        }();
        return ops;
    }

    CrystalLattice ApplyBasisOp(const CrystalLattice &in, const BasisOp &op) {
        const Coord v[3] = {in.Vec0(), in.Vec1(), in.Vec2()};
        return CrystalLattice(v[op.perm[0]] * static_cast<float>(op.signs[0]),
                              v[op.perm[1]] * static_cast<float>(op.signs[1]),
                              v[op.perm[2]] * static_cast<float>(op.signs[2]));
    }
}

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

    for (size_t i = 1; i < lattices.size(); i++) {
        for (const auto &op : RightHandedOps()) {
            const CrystalLattice latt = ApplyBasisOp(lattices[i], op);

            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();
            ret.push_back({Coord(static_cast<float>(rod.x()),
                                 static_cast<float>(rod.y()),
                                 static_cast<float>(rod.z()))});
            break;
        }
    }

    return ret;
}