#!/usr/bin/env python3
"""Three facts about high-dimensional geometry, checked numerically.

1. Random directions are nearly orthogonal.  For independent uniformly random
   unit vectors u, v in R^d, cos(u, v) has mean 0 and standard deviation
   exactly 1/sqrt(d); it concentrates like a N(0, 1/d) variable.
2. Almost all volume of a ball is near its surface.  The fraction of the unit
   d-ball within distance eps of the boundary is 1 - (1 - eps)^d -> 1.
3. Many more than d 'almost orthogonal' directions fit.  Only d vectors can
   be exactly orthogonal, but the number of unit vectors with all pairwise
   |cos| < 0.1 grows exponentially in d.  Two views:
   (a) purely random sets: P(|cos| >= 0.1) for one random pair is
       p_d = P(Beta(1/2, (d-1)/2) >= 0.01) ~ exp(-d/200); by the birthday
       argument about N ~ sqrt(2 ln 2 / p_d) random vectors survive with
       probability 1/2 -- tiny for d = 1000 (~30), huge for d = 4000 (~10^5);
   (b) greedy search (keep a random candidate only if it is compatible with
       all kept ones) does better: with 20k tries it passes d between
       d = 400 and d = 800.
   This is why embeddings can store many more 'features' than dimensions.

  python high_dim.py [--quick]
Writes out/high_dim.png and prints the numbers.
"""
from __future__ import annotations

import argparse
import os

import numpy as np

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out")


def random_unit(n, d, rng):
    v = rng.standard_normal((n, d))
    return v / np.linalg.norm(v, axis=1, keepdims=True)


def cosine_samples(d, n_pairs=20000, seed=0):
    """Cosines between n_pairs independent pairs of random unit vectors."""
    rng = np.random.default_rng(seed)
    out = np.empty(n_pairs)
    batch = max(1, min(n_pairs, 2_000_000 // d))
    for s in range(0, n_pairs, batch):
        m = min(batch, n_pairs - s)
        u, v = random_unit(m, d, rng), random_unit(m, d, rng)
        out[s:s + m] = (u * v).sum(axis=1)
    return out


def shell_fraction(d, eps):
    """Fraction of the unit d-ball's volume at distance < eps from the surface."""
    return 1.0 - (1.0 - eps) ** d


def shell_fraction_mc(d, eps, n=200000, seed=0):
    """Monte Carlo check by rejection sampling from the cube (small d only)."""
    rng = np.random.default_rng(seed)
    x = rng.uniform(-1, 1, (n, d))
    r = np.linalg.norm(x, axis=1)
    inside = r <= 1
    return float(np.mean(r[inside] > 1 - eps)), int(inside.sum())


def greedy_packing(d, thresh=0.1, attempts=20000, batch=256, seed=0):
    """Greedily keep random unit vectors whose |cos| with all kept ones < thresh.

    Returns the number kept after `attempts` random candidates (a lower bound
    on the maximum size of such a set)."""
    rng = np.random.default_rng(seed)
    kept = np.empty((0, d))
    tried = 0
    while tried < attempts:
        m = min(batch, attempts - tried)
        cand = random_unit(m, d, rng)
        tried += m
        if len(kept):
            ok = np.all(np.abs(cand @ kept.T) < thresh, axis=1)
            cand = cand[ok]
        # resolve conflicts inside the batch greedily, in order
        chosen = []
        for i in range(len(cand)):
            if all(abs(cand[i] @ cand[j]) < thresh for j in chosen):
                chosen.append(i)
        if chosen:
            kept = np.vstack([kept, cand[chosen]])
    return len(kept)


def pair_violation_prob(d, thresh=0.1):
    """Exact P(|cos| >= thresh) for two random unit vectors in R^d:
    cos^2 ~ Beta(1/2, (d-1)/2)."""
    from scipy import stats
    return float(stats.beta.sf(thresh**2, 0.5, (d - 1) / 2))


def random_set_size(d, thresh=0.1, trials=5, seed=0):
    """Median number of i.i.d. random unit vectors drawn before the first
    pair with |cos| >= thresh appears."""
    rng = np.random.default_rng(seed)
    p = pair_violation_prob(d, thresh)
    n_max = int(min(8000, max(50, 4 * np.sqrt(2 * np.log(2) / p))))
    sizes = []
    for _ in range(trials):
        V = random_unit(n_max, d, rng).astype(np.float32)
        G = np.abs(V @ V.T) >= thresh
        np.fill_diagonal(G, False)
        G = np.triu(G)                       # G[i, j] with i < j
        bad_cols = np.nonzero(G.any(axis=0))[0]
        sizes.append(int(bad_cols[0]) if len(bad_cols) else n_max)  # vectors 0..j-1 are fine
    return float(np.median(sizes))


def main(quick=False):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    os.makedirs(OUT, exist_ok=True)
    dims = [2, 3, 10, 100, 1000, 10000]
    print("1) cosine similarity of random unit vectors")
    print(f"{'d':>7} {'std(cos)':>10} {'1/sqrt(d)':>10} {'P(|cos|>0.1)':>13}")
    cos = {}
    for d in dims:
        c = cosine_samples(d, n_pairs=5000 if quick else 20000, seed=d)
        cos[d] = c
        print(f"{d:>7} {c.std():>10.4f} {1 / np.sqrt(d):>10.4f} {np.mean(np.abs(c) > 0.1):>13.3f}")

    print("\n2) fraction of ball volume within eps of the surface: 1-(1-eps)^d")
    for d in (2, 3, 5, 8):
        mc, n_in = shell_fraction_mc(d, 0.1, n=50000 if quick else 200000, seed=d)
        print(f"   d={d:<3} eps=0.1  formula {shell_fraction(d, 0.1):.4f}   Monte Carlo {mc:.4f} ({n_in} pts)")
    for d in (10, 100, 1000):
        print(f"   d={d:<5} eps=0.01 {shell_fraction(d, 0.01):.4f}   eps=0.1 {shell_fraction(d, 0.1):.6f}")

    print("\n3) greedy random packing with all pairwise |cos| < 0.1")
    pack_dims = [10, 50, 100, 200, 400] if quick else [10, 50, 100, 200, 400, 800]
    attempts = 4000 if quick else 20000
    packs = []
    for d in pack_dims:
        k = greedy_packing(d, 0.1, attempts=attempts, seed=d)
        packs.append(k)
        print(f"   d={d:<4} kept {k:>5} vectors  ({k / d:.2f} x d)  from {attempts} random tries")

    print("\n   random sets (no search): median # vectors before the first |cos| >= 0.1")
    rand_dims = [200, 500, 1000, 1500] if quick else [200, 500, 1000, 1500, 2000, 2500]
    rand_sizes = []
    for d in rand_dims:
        n_sim = random_set_size(d, trials=3 if quick else 5, seed=d)
        rand_sizes.append(n_sim)
        p = pair_violation_prob(d)
        print(f"   d={d:<5} simulated {n_sim:>7.0f}   birthday estimate sqrt(2 ln2/p) = "
              f"{np.sqrt(2 * np.log(2) / p):>8.0f}   (p = {p:.2e})")
    for d in (3000, 4000, 5000):
        p = pair_violation_prob(d)
        print(f"   d={d:<5} birthday estimate {np.sqrt(2 * np.log(2) / p):.3g}   (p = {p:.2e})")

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.4))
    ax = axes[0]
    bins = np.linspace(-1, 1, 201)
    for d in dims:
        ax.hist(cos[d], bins=bins, density=True, histtype="step", lw=1.4, label=f"d={d}")
    ax.set_yscale("log")
    ax.set_ylim(1e-2, None)
    ax.set_xlabel("cos(u, v) for random unit u, v")
    ax.set_ylabel("density")
    ax.set_title("Random directions are nearly orthogonal\n(std = 1/√d)")
    ax.legend(fontsize=8)

    ax = axes[1]
    ds = np.arange(1, 1001)
    for eps in (0.01, 0.05, 0.1):
        ax.plot(ds, shell_fraction(ds, eps), label=f"outer {eps:.0%} of radius")
    ax.set_xscale("log")
    ax.set_xlabel("dimension d")
    ax.set_ylabel("fraction of ball volume")
    ax.set_title("Volume concentrates near the surface")
    ax.legend(fontsize=8)

    ax = axes[2]
    ax.plot(pack_dims, packs, "o-", label="greedy search (kept / 20k tries)" if not quick
            else "greedy search")
    ax.plot(rand_dims, rand_sizes, "s", color="C1", label="i.i.d. random set (simulated)")
    dd = np.arange(100, 5001, 50)
    ax.plot(dd, [np.sqrt(2 * np.log(2) / pair_violation_prob(d)) for d in dd], "-", color="C1",
            lw=1, label="i.i.d. random set (birthday estimate)")
    ax.plot(dd, dd, "k--", label="d (exactly orthogonal max)")
    ax.set_yscale("log")
    ax.set_xlabel("dimension d")
    ax.set_ylabel("# unit vectors, all |cos| < 0.1")
    ax.set_title("Many more 'almost orthogonal' directions than d")
    ax.legend(fontsize=8)
    for a in axes:
        a.grid(alpha=0.3)
    fig.tight_layout()
    path = os.path.join(OUT, "high_dim.png")
    fig.savefig(path, dpi=130)
    plt.close(fig)
    print(f"saved {path}")
    return {"pack_dims": pack_dims, "packs": packs}


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--quick", action="store_true")
    main(ap.parse_args().quick)