"""Claim 2 verification: Proposition 5.1 & Corollary 5.2.

Prop 5.1: p_hat_{mu,K}(x_i) = 1/(N V_d mu_i^d) * (1/K Sum_k k^{1/d})^d is a local
density estimator (asymptotically p(x_i)).
Cor 5.2: at ICDM convergence, mu_i ~ mu_bar for all i => p_hat equal everywhere
(density uniformized).

We verify numerically:
 (a) p_hat_{mu,K} correlates with the true density in a mixture of Gaussians.
 (b) after ICDM, the spread of mu_i collapses (std -> ~0) while raw std is large,
     confirming density gradient removal.
"""
import numpy as np
import sys, os
sys.path.insert(0, os.path.dirname(__file__))
from gicdm_core import pairwise_sq_dists, icdm_scaling, knn_distances


def unit_ball_volume(d):
    from math import gamma, pi
    return (pi ** (d / 2)) / gamma(d / 2 + 1)


def p_hat_mu_k(X, K):
    D = pairwise_sq_dists(X)
    mu = knn_distances(D, K, self_included=True).mean(axis=1)
    N, d = X.shape
    Vd = unit_ball_volume(d)
    w = (np.sum([k ** (1.0 / d) for k in range(1, K + 1)]) / K) ** d
    return 1.0 / (N * Vd * mu ** d) * w, mu


def run():
    rng = np.random.default_rng(0)
    # Mixture of two Gaussians with different variances -> different densities
    d = 20
    n1, n2 = 800, 800
    X = np.vstack([
        rng.normal(0, 0.5, size=(n1, d)),
        rng.normal(3, 2.0, size=(n2, d)),
    ])
    true_dens = np.concatenate([
        (2 * np.pi * 0.25) ** (-d / 2) * np.ones(n1),
        (2 * np.pi * 4.0) ** (-d / 2) * np.ones(n2),
    ])
    K = 20
    phat, mu = p_hat_mu_k(X, K)
    # correlation between estimator and true density (log scale, robust)
    log_corr = np.corrcoef(np.log(phat), np.log(true_dens))[0, 1]

    # ICDM convergence: spread of mu before/after
    mu_before = mu.copy()
    delta, mu_after = icdm_scaling(pairwise_sq_dists(X), K, n_iter=10, return_mu=True)
    res = dict(
        d=d, N=len(X), K=K,
        log_corr_density=float(log_corr),
        mu_before_std=float(mu_before.std()),
        mu_before_cv=float(mu_before.std() / mu_before.mean()),
        mu_after_std=float(mu_after.std()),
        mu_after_cv=float(mu_after.std() / mu_after.mean()),
    )
    print("Prop 5.1 log-correlation p_hat vs true density:", round(res['log_corr_density'], 3))
    print("mu_i std before ICDM:", round(res['mu_before_std'], 4),
          "| after ICDM:", round(res['mu_after_std'], 5))
    print("mu_i CV  before ICDM:", round(res['mu_before_cv'], 4),
          "| after ICDM:", round(res['mu_after_cv'], 5))
    return res


if __name__ == "__main__":
    import json
    res = run()
    json.dump(res, open("results/claim2_density.json", "w"))