import { assert } from "@woosh/meep-engine/src/core/assert.js";

let scratch_next = new Float64Array(16);

function ensure_scratch(n) {
    if (scratch_next.length < n) scratch_next = new Float64Array(n);
    return scratch_next;
}

/**
 * Find the eigenvector corresponding to the largest-magnitude eigenvalue of a
 * symmetric square matrix using power iteration.
 *
 * Designed for small symmetric matrices where the dominant eigenvalue is
 * well-separated from the rest (Horn's quaternion method on a 4x4, PCA on a
 * 3x3 covariance, etc.). Unsuitable when the top two eigenvalues have nearly
 * equal magnitude — convergence stalls in that regime.
 *
 * @param {SquareMatrix} matrix Symmetric, column-major as per SquareMatrix
 * @param {number[]|Float32Array|Float64Array} out Eigenvector is written here, length >= matrix.size
 * @param {number} [max_iterations=64]
 * @param {number} [tolerance=1e-8] Convergence stops when ||v_next - v|| < tolerance
 * @returns {number} The Rayleigh quotient (estimated eigenvalue)
 */
export function matrix_top_eigenvector_power_iteration(
    matrix,
    out,
    max_iterations = 64,
    tolerance = 1e-8
) {
    assert.defined(matrix, "matrix");
    assert.defined(out, "out");

    const n = matrix.size;
    const data = matrix.data;

    // Deterministic asymmetric initial vector. A uniform init (e.g. 1/√n in
    // every slot) can be orthogonal to the dominant eigenvector when the matrix
    // has block structure, stalling iteration on a sub-dominant eigenvalue.
    // Irrational-component cosines avoid this for any reasonable n.
    let init_norm_sq = 0;
    for (let i = 0; i < n; i++) {
        const v = Math.cos(i * 1.7 + 0.3);
        out[i] = v;
        init_norm_sq += v * v;
    }
    const init_inv = 1 / Math.sqrt(init_norm_sq);
    for (let i = 0; i < n; i++) out[i] *= init_inv;

    const next = ensure_scratch(n);
    let eigenvalue = 0;

    for (let iter = 0; iter < max_iterations; iter++) {
        // next = matrix * out
        // Column-major: data[col * n + row] is element at (row, col)
        for (let row = 0; row < n; row++) {
            let sum = 0;
            for (let col = 0; col < n; col++) {
                sum += data[col * n + row] * out[col];
            }
            next[row] = sum;
        }

        // Rayleigh quotient = out . next  (since out is unit length and matrix is symmetric)
        let rayleigh = 0;
        for (let i = 0; i < n; i++) {
            rayleigh += out[i] * next[i];
        }

        // Normalize next
        let length_sqr = 0;
        for (let i = 0; i < n; i++) {
            length_sqr += next[i] * next[i];
        }

        if (length_sqr === 0) {
            // Degenerate input — bail out, leave `out` as-is
            return 0;
        }

        const inv_length = 1 / Math.sqrt(length_sqr);

        // Sign-align next with out so we can measure convergence by raw difference.
        // (Power iteration has a sign ambiguity per step when eigenvalue is negative.)
        let dot = 0;
        for (let i = 0; i < n; i++) {
            dot += next[i] * out[i];
        }
        const sign = dot < 0 ? -inv_length : inv_length;

        let delta_sqr = 0;
        for (let i = 0; i < n; i++) {
            const v = next[i] * sign;
            const d = v - out[i];
            delta_sqr += d * d;
            out[i] = v;
        }

        eigenvalue = rayleigh;

        if (delta_sqr < tolerance * tolerance) {
            break;
        }
    }

    return eigenvalue;
}
