"""How much does the 2-D picture lie? Two standard checks of a layout against the original space.

* Trustworthiness (Venna & Kaski, 2001): 1 means every point's k nearest neighbours on the map were
  also among its nearest neighbours in the original space; it falls when the map puts strangers next to
  each other, weighted by how far away they really were (same definition as scikit-learn's).
* Neighbourhood preservation: the plain fraction of each point's k original neighbours that are still
  among its k map neighbours, averaged over points.
"""
from __future__ import annotations

import numpy as np


def pairwise_sq(X: np.ndarray) -> np.ndarray:
    diff = X[:, None, :] - X[None, :, :]
    return np.einsum("ijk,ijk->ij", diff, diff)


def knn(D: np.ndarray, k: int) -> np.ndarray:
    """Indices of the k nearest other points for every row of a distance matrix (ties by index)."""
    n = D.shape[0]
    out = np.zeros((n, k), dtype=int)
    for i in range(n):
        order = sorted((float(D[i, j]), j) for j in range(n) if j != i)
        out[i] = [j for _, j in order[:k]]
    return out


def trustworthiness(X_high: np.ndarray, Y_low: np.ndarray, k: int = 5) -> float:
    n = X_high.shape[0]
    if n <= k + 1:
        return 1.0
    Dh, Dl = pairwise_sq(X_high), pairwise_sq(Y_low)
    nn_low = knn(Dl, k)
    total = 0.0
    for i in range(n):
        order = sorted((float(Dh[i, j]), j) for j in range(n) if j != i)
        rank = {j: r + 1 for r, (_, j) in enumerate(order)}      # 1 = nearest in the original space
        for j in nn_low[i]:
            if rank[j] > k:
                total += rank[j] - k
    return 1.0 - 2.0 / (n * k * (2 * n - 3 * k - 1)) * total


def neighbourhood_preservation(X_high: np.ndarray, Y_low: np.ndarray, k: int = 5) -> float:
    n = X_high.shape[0]
    if n <= k:
        return 1.0
    nh, nl = knn(pairwise_sq(X_high), k), knn(pairwise_sq(Y_low), k)
    return float(np.mean([len(set(nh[i]) & set(nl[i])) / k for i in range(n)]))