"""Distributional metrics in state space.

``energy_distance`` is the validation metric of the paper's model selection (a V-statistic: ``md(X, X)`` and
``md(Y, Y)`` average over all pairs including self-pairs; ``2 md(X, Y) - md(X, X) - md(Y, Y)``).
"""
from __future__ import annotations

import numpy as np


def _mean_pairwise_dist(A: np.ndarray, B: np.ndarray) -> float:
    return np.mean(np.linalg.norm(A[:, None] - B[None, :], axis=-1))


def energy_distance(X: np.ndarray, Y: np.ndarray) -> float:
    md = _mean_pairwise_dist
    return float(2 * md(X, Y) - md(X, X) - md(Y, Y))


def energy_distance_cached_self(X: np.ndarray, Y: np.ndarray, md_yy: float) -> float:
    """``energy_distance`` with the truth-self term ``md(Y, Y)`` supplied (it does not depend on X)."""
    md = _mean_pairwise_dist
    return float(2 * md(X, Y) - md(X, X) - float(md_yy))


def energy_distance_u(X: np.ndarray, Y: np.ndarray) -> float:
    """The unbiased U-statistic: the self terms exclude the diagonal, ``ED_U = ED_V - md(X,X)/(M-1) - md(Y,Y)/(N-1)``."""
    md = _mean_pairwise_dist
    mxx, myy = md(X, X), md(Y, Y)
    M, N = X.shape[0], Y.shape[0]
    return float(2 * md(X, Y) - mxx - myy - (mxx / (M - 1) if M > 1 else 0.0) - (myy / (N - 1) if N > 1 else 0.0))


def effective_rank(X: np.ndarray) -> float:
    """Effective rank of the sample covariance, tr(S)^2 / ||S||_F^2; NaN for fewer than 2 samples."""
    X = np.asarray(X, dtype=np.float64)
    if X.shape[0] < 2:
        return float("nan")
    S = np.cov(X.T)
    fro2 = float((S * S).sum())
    return float(S.trace() ** 2 / fro2) if fro2 > 0 else float("nan")