"""A compact exact t-SNE for n <= 500 points (van der Maaten & Hinton, 2008), written to match web/map.js
step for step. No Barnes-Hut tree: the room is small enough that the plain O(n^2) gradient is fast.

What it does, in plain words: for every pair of phrases it works out how likely one would be picked as
a "neighbour" of the other in the 10-dimensional PCA space (the P matrix, using a Gaussian per point whose
width is tuned so that each point has about `perplexity` effective neighbours). It then moves points
around on a 2-D sheet so that the same kind of neighbour probabilities on the sheet (the Q matrix, using
a heavy-tailed Student-t kernel) match P as closely as possible, by gradient descent on KL(P || Q).

Deterministic by construction: the starting positions are the first two PCA components (scaled tiny),
not random numbers, so the same room always gives the same map.
"""
from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np

EARLY_ITERS = 100          # iterations with exaggeration 12 and momentum 0.5
EXAGGERATION = 12.0
ITERS = 500
MIN_GAIN = 0.01


def squared_distances(X: np.ndarray) -> np.ndarray:
    """Pairwise squared Euclidean distances, by direct differences (the expanded |a|^2 + |b|^2 - 2ab form
    loses digits for near-identical rows, and the JavaScript uses differences too)."""
    diff = X[:, None, :] - X[None, :, :]
    return np.einsum("ijk,ijk->ij", diff, diff)


def affinities(D2: np.ndarray, perplexity: float, tol: float = 1e-5, iters: int = 50) -> np.ndarray:
    """Symmetric joint probabilities P from squared distances, one Gaussian width per point found by
    bisection so that the entropy of its row equals log(perplexity)."""
    n = D2.shape[0]
    P = np.zeros((n, n))
    logU = np.log(perplexity)
    for i in range(n):
        idx = np.concatenate([np.arange(i), np.arange(i + 1, n)])
        Di = D2[i, idx]
        beta, lo, hi = 1.0, -np.inf, np.inf
        for _ in range(iters):
            Pi = np.exp(-Di * beta)
            sumP = max(float(np.sum(Pi)), 1e-12)
            H = np.log(sumP) + beta * float(np.sum(Di * Pi)) / sumP
            Hdiff = H - logU
            if abs(Hdiff) < tol:
                break
            if Hdiff > 0:
                lo = beta
                beta = beta * 2.0 if hi == np.inf else (beta + hi) / 2.0
            else:
                hi = beta
                beta = beta / 2.0 if lo == -np.inf else (beta + lo) / 2.0
        P[i, idx] = Pi / sumP
    P = (P + P.T) / (2.0 * n)
    return np.maximum(P, 1e-12)


def default_perplexity(n: int) -> float:
    return max(2.0, min(30.0, (n - 1) / 3.0))


def initial_positions(pcs: np.ndarray) -> np.ndarray:
    """First two PCA components, scaled so the first has standard deviation 1e-4 (Kobak & Berens, 2019)."""
    n = pcs.shape[0]
    Y = np.zeros((n, 2))
    k = min(2, pcs.shape[1])
    Y[:, :k] = pcs[:, :k]
    sd = float(np.std(Y[:, 0]))
    if sd > 0:
        Y *= 1e-4 / sd
    return Y


@dataclass
class TSNE:
    """Gradient-descent state, so a layout can be advanced a few iterations at a time (and animated)."""
    P: np.ndarray
    Y: np.ndarray
    iY: np.ndarray = field(init=False)
    gains: np.ndarray = field(init=False)
    it: int = 0
    eta: float = 0.0

    def __post_init__(self) -> None:
        n = self.Y.shape[0]
        self.iY = np.zeros_like(self.Y)
        self.gains = np.ones_like(self.Y)
        self.eta = max(n / 48.0, 50.0)       # scikit-learn's "auto" learning rate

    def q_and_num(self) -> tuple[np.ndarray, np.ndarray]:
        num = 1.0 / (1.0 + squared_distances(self.Y))
        np.fill_diagonal(num, 0.0)
        Q = np.maximum(num / num.sum(), 1e-12)
        return Q, num

    def step(self, iters: int = 1) -> None:
        n = self.Y.shape[0]
        for _ in range(iters):
            early = self.it < EARLY_ITERS
            Peff = self.P * EXAGGERATION if early else self.P
            momentum = 0.5 if early else 0.8
            Q, num = self.q_and_num()
            PQ = (Peff - Q) * num
            # dY_i = 4 * sum_j PQ_ij (y_i - y_j)
            dY = 4.0 * (np.diag(PQ.sum(axis=1)) @ self.Y - PQ @ self.Y)
            same = np.sign(dY) == np.sign(self.iY)
            self.gains = np.where(same, self.gains * 0.8, self.gains + 0.2)
            self.gains = np.maximum(self.gains, MIN_GAIN)
            self.iY = momentum * self.iY - self.eta * (self.gains * dY)
            self.Y = self.Y + self.iY
            self.Y -= self.Y.mean(axis=0)
            self.it += 1

    def kl(self) -> float:
        Q, _ = self.q_and_num()
        P = self.P
        return float(np.sum(P * np.log(P / Q)))


def layout(pcs: np.ndarray, perplexity: float | None = None, iters: int = ITERS) -> tuple[np.ndarray, float]:
    """Full layout of the PCA scores. Returns (Y, KL divergence). Tiny rooms (n < 5) just use PCA."""
    n = pcs.shape[0]
    if n < 5:
        Y = initial_positions(pcs) * 1e4
        return Y - (Y.mean(axis=0) if n else 0), 0.0
    P = affinities(squared_distances(pcs), perplexity or default_perplexity(n))
    t = TSNE(P, initial_positions(pcs))
    t.step(iters)
    return t.Y, t.kl()