"""Proposition 5.3: crossover dimension d* for standard Gaussian N(0, I_d).

d* solves  P( median NND_k^2 <= E||X||^2 ) = 1/2, i.e. the integral
  int_0^infty P( Bin(N-1, F_{chi2_d}(lambda=r)) >= k ) f_{chi2_d}(r) dr = 1/2
where F_{chi2_d}(lambda) is the noncentral chi2 CDF with noncentrality lambda,
f_{chi2_d} is the chi2 density (central, lambda=0), and r = ||X||^2 ~ chi2_d.
"""
import numpy as np
from scipy import integrate, stats


def crossover_dimension(N, k, d_grid=None):
    """Numerically solve Proposition 5.3 for d* given N and k."""
    if d_grid is None:
        d_grid = np.arange(2, 4001)  # up to 4000 dims

    def median_prob(d):
        # P( Bin(N-1, F_{chi2_d}(lambda=r)) >= k ) averaged over r ~ chi2_d
        def integrand(r):
            # noncentral chi2 cdf with noncentrality lambda = r
            # F_{chi2_d}(lambda=r)(d) = P( chi2_d(r) <= d )
            p = stats.ncx2.cdf(d, d, r)
            return stats.binom.cdf(k - 1, N - 1, p, loc=1)  # P(Bin >= k) = 1 - P(Bin <= k-1)
        # integrate r over chi2_d density
        val, _ = integrate.quad(lambda r: stats.chi2.pdf(r, d) * (1 - stats.binom.cdf(k - 1, N - 1, stats.ncx2.cdf(d, d, r))),
                                0, 200 + 20 * d, limit=200)
        return val

    probs = []
    for d in d_grid:
        probs.append(median_prob(d))
    probs = np.array(probs)
    # find d* where prob crosses 1/2
    idx = np.argmin(np.abs(probs - 0.5))
    return d_grid[idx], probs


if __name__ == "__main__":
    for N in [1000, 10000, 50000]:
        for k in [1, 5, 10, 20]:
            dstar, _ = crossover_dimension(N, k)
            print(f"N={N:6d} k={k:3d} -> d*={dstar}")