import { BinaryDataType } from "../../binary/type/BinaryDataType.js";
import { matrix_top_eigenvector_power_iteration } from "../../math/linalg/eigen/matrix_top_eigenvector_power_iteration.js";
import { SquareMatrix } from "../../math/matrix/SquareMatrix.js"

const scratch_horn_n = new SquareMatrix(4, BinaryDataType.Float64);

/**
 * Closed-form solution for the rotation that best aligns two sets of unit
 * vectors, via Horn (1987) "Closed-form solution of absolute orientation using
 * unit quaternions".
 *
 * Given paired unit vectors {a_i} and {b_i}, finds the unit quaternion q such
 * that rotating each b_i by q minimizes Σ |a_i − R(q)·b_i|².
 *
 * The matrix N is built so that the eigenvector of its largest eigenvalue
 * IS the optimal rotation quaternion (in scalar-last form, matching meep's
 * `Quaternion`).
 *
 * Inputs are unit vectors so no centroid subtraction is needed — the rotation
 * is recovered from the cross-correlations alone.
 *
 * @param {Float32Array|number[]} a_xyz Source vectors — Length: count*3, unit length per triple
 * @param {Float32Array|number[]} b_xyz Target vectors — Length: count*3, unit length per triple
 * @param {number} count
 * @param {Float32Array|number[]} out_quaternion Length 4 (x, y, z, w)
 * @returns {number} The maximum eigenvalue ≈ Σ a_i · R·b_i. For a perfect alignment this equals `count`.
 */
export function v3_rigid_align_paired_unit_vectors(a_xyz, b_xyz, count, out_quaternion) {

    // Cross-correlation matrix S_ij = Σ b_i_component * a_j_component
    // Note: S is built as Σ b ⊗ aᵀ (Horn 1987 convention) — the rotation
    // recovered will take b onto a.
    let Sxx = 0, Sxy = 0, Sxz = 0;
    let Syx = 0, Syy = 0, Syz = 0;
    let Szx = 0, Szy = 0, Szz = 0;

    for (let i = 0; i < count; i++) {
        const i3 = i * 3;
        const ax = a_xyz[i3], ay = a_xyz[i3 + 1], az = a_xyz[i3 + 2];
        const bx = b_xyz[i3], by = b_xyz[i3 + 1], bz = b_xyz[i3 + 2];

        Sxx += bx * ax;
        Sxy += bx * ay;
        Sxz += bx * az;
        Syx += by * ax;
        Syy += by * ay;
        Syz += by * az;
        Szx += bz * ax;
        Szy += bz * ay;
        Szz += bz * az;
    }

    // Horn's 4x4 symmetric matrix N (scalar-last quaternion ordering: x, y, z, w).
    //
    // From the paper, ordered for scalar-FIRST quaternion (w, x, y, z):
    //   N00 = Sxx + Syy + Szz       N01 = Syz - Szy            N02 = Szx - Sxz            N03 = Sxy - Syx
    //   N11 = Sxx - Syy - Szz       N12 = Sxy + Syx            N13 = Szx + Sxz
    //   N22 = -Sxx + Syy - Szz      N23 = Syz + Szy
    //   N33 = -Sxx - Syy + Szz
    //
    // We rebind to (x, y, z, w):
    //   row/col 0 = x, 1 = y, 2 = z, 3 = w
    // i.e. swap original index 0 with 3 and shift others down.
    const Nxx = Sxx - Syy - Szz;
    const Nyy = -Sxx + Syy - Szz;
    const Nzz = -Sxx - Syy + Szz;
    const Nww = Sxx + Syy + Szz;
    const Nxy = Sxy + Syx;
    const Nxz = Szx + Sxz;
    const Nxw = Syz - Szy;
    const Nyz = Syz + Szy;
    const Nyw = Szx - Sxz;
    const Nzw = Sxy - Syx;

    const N = scratch_horn_n;
    const d = N.data;
    // SquareMatrix is column-major: d[col * 4 + row]
    d[0] = Nxx;
    d[5] = Nyy;
    d[10] = Nzz;
    d[15] = Nww;
    d[1] = Nxy;
    d[4] = Nxy;
    d[2] = Nxz;
    d[8] = Nxz;
    d[3] = Nxw;
    d[12] = Nxw;
    d[6] = Nyz;
    d[9] = Nyz;
    d[7] = Nyw;
    d[13] = Nyw;
    d[11] = Nzw;
    d[14] = Nzw;

    return matrix_top_eigenvector_power_iteration(N, out_quaternion);
}
