"""PCA: the first step of the layout. Embeddings have 384 or 1024 numbers each; t-SNE only needs the
directions along which the room actually varies, so we keep the top 10 principal components.

The embeddings are unit vectors (cosine-normalised), so Euclidean distance between them is a monotone
function of cosine similarity: ||a - b||^2 = 2 - 2 cos(a, b). PCA on them then preserves exactly the
geometry t-SNE will look at.

Because there are far fewer phrases than dimensions (n <= 500, d >= 384), the eigenproblem is solved on
the n x n Gram matrix instead of the d x d covariance: X X^T has the same non-zero eigenvalues as X^T X,
and the scores are just its eigenvectors scaled by the square roots of the eigenvalues.
"""
from __future__ import annotations

import numpy as np


def normalise_rows(X: np.ndarray) -> np.ndarray:
    """Scale every row to unit length (rows of zeros stay zero)."""
    X = np.asarray(X, dtype=float)
    n = np.linalg.norm(X, axis=1, keepdims=True)
    return np.where(n > 0, X / np.where(n > 0, n, 1.0), 0.0)


def pca(X: np.ndarray, k: int = 10) -> tuple[np.ndarray, np.ndarray]:
    """Project the rows of X onto their top-k principal components.

    Returns (scores, explained) where scores is n x k' (k' = min(k, n-1, d)) and explained is the
    fraction of total variance each component carries. Signs of components are made deterministic
    (largest-magnitude entry of each score column positive) so plots are reproducible, but any
    downstream use must not depend on the sign, since the JavaScript may choose the other.
    """
    X = np.asarray(X, dtype=float)
    n, d = X.shape
    k = max(0, min(k, n - 1, d))
    if n == 0 or k == 0:
        return np.zeros((n, 0)), np.zeros(0)
    Xc = X - X.mean(axis=0)
    G = Xc @ Xc.T                           # n x n Gram matrix
    evals, evecs = np.linalg.eigh(G)        # ascending
    order = np.argsort(evals)[::-1][:k]
    lam = np.clip(evals[order], 0.0, None)
    U = evecs[:, order]
    scores = U * np.sqrt(lam)               # = Xc @ components^T
    total = float(np.trace(G))
    explained = lam / total if total > 0 else np.zeros(k)
    # deterministic signs
    for j in range(scores.shape[1]):
        i = int(np.argmax(np.abs(scores[:, j])))
        if scores[i, j] < 0:
            scores[:, j] = -scores[:, j]
    return scores, explained