import json
import math
import time
from pathlib import Path

import numpy as np
from scipy.integrate import quad
from scipy.stats import beta

SEED = 3044
rng = np.random.default_rng(SEED)


def exact_tau(m, d):
    """Theorem 7.1: expected Kendall correlation for one distance comparison."""
    f = lambda b: math.asin(math.sqrt(b)) * beta.pdf(b, m / 2.0, (d - m) / 2.0)
    return 2.0 / math.pi * quad(f, 0.0, 1.0, epsabs=2e-8)[0]


def asymptotic_tau(m, d):
    return 2.0 / math.pi * math.sqrt(m / d)


def make_projection(m, d, seed):
    # Haar random rank-m projector represented by an orthonormal-row matrix.
    local = np.random.default_rng(seed)
    q, _ = np.linalg.qr(local.normal(size=(d, m)), mode="reduced")
    return q.T.astype(np.float32)


def pairwise_order_check(d, m, n_pairs=30000, seed=0):
    local = np.random.default_rng(seed)
    r = make_projection(m, d, seed + 10000)
    # Independent triples are the exact theorem setup.
    x0 = local.normal(size=(n_pairs, d)).astype(np.float32)
    x1 = local.normal(size=(n_pairs, d)).astype(np.float32)
    x2 = local.normal(size=(n_pairs, d)).astype(np.float32)
    full = ((x1 - x0) ** 2).sum(1) - ((x2 - x0) ** 2).sum(1)
    y0, y1, y2 = x0 @ r.T, x1 @ r.T, x2 @ r.T
    proj = ((y1 - y0) ** 2).sum(1) - ((y2 - y0) ** 2).sum(1)
    agreement = np.mean(np.sign(full) == np.sign(proj))
    tau = 2 * agreement - 1
    return float(tau), float(agreement)


def recall_at(retrieved, truth, k):
    return float(np.mean([len(set(a[:k]) & set(b[:k])) / k for a, b in zip(retrieved, truth)]))


def adaptive_dimension(d, target_tau=0.5, safety_margin=0.0, max_rounds=8):
    """Select m by the proposed formula, then expand using held-out pairwise tau."""
    m = max(8, int(math.ceil(d * (math.pi * target_tau / 2) ** 2)))
    m = min(d, int(math.ceil(m / 8) * 8))
    history = []
    for round_id in range(max_rounds):
        tau, _ = pairwise_order_check(d, m, n_pairs=12000,
                                      seed=SEED + 5000 + round_id)
        history.append({"round": round_id, "m": m, "measured_tau": tau})
        if tau >= target_tau - safety_margin or m >= d:
            break
        m = min(d, max(m + 1, int(math.ceil(1.25 * m / 8) * 8)))
    return m, history

def retrieval_trial(d=128, n=4000, nq=300, target_tau=0.5):
    local = np.random.default_rng(SEED + 77)
    # A mildly clustered retrieval-like data set: exact neighbors share a latent center.
    n_centers = 80
    centers = local.normal(size=(n_centers, d)).astype(np.float32)
    labels = local.integers(0, n_centers, size=n)
    data = centers[labels] + 0.8 * local.normal(size=(n, d)).astype(np.float32)
    qlabels = local.integers(0, n_centers, size=nq)
    queries = centers[qlabels] + 0.8 * local.normal(size=(nq, d)).astype(np.float32)
    # Full-space truth and candidate pair ranking metric.
    full_dist = ((queries[:, None, :] - data[None, :, :]) ** 2).sum(2)
    truth = np.argsort(full_dist, axis=1)
    # Standard conservative JL sizing (one common constant); cap at d because no
    # projection can improve on the original representation.
    eps = 0.5
    jl_m = min(d, int(math.ceil(8 * math.log(n) / (eps * eps))))
    rank_m = max(8, int(math.ceil(d * (math.pi * target_tau / 2) ** 2)))
    # Hardware-friendly rounding, while retaining the requested target budget.
    rank_m = min(d, int(math.ceil(rank_m / 8) * 8))
    rows = {"jl": jl_m, "ranking_rule": rank_m}
    out = {}
    for name, m in rows.items():
        t0 = time.perf_counter()
        r = make_projection(m, d, SEED + m)
        zd, zq = data @ r.T, queries @ r.T
        pd = ((zq[:, None, :] - zd[None, :, :]) ** 2).sum(2)
        got = np.argsort(pd, axis=1)
        # Kendall tau over the top-20 candidate ordering, measured against full space.
        taus = []
        for i in range(nq):
            ids = truth[i, :20]
            a = full_dist[i, ids]
            b = pd[i, ids]
            signs = np.sign((a[:, None] - a[None, :]) * (b[:, None] - b[None, :]))
            iu = np.triu_indices(len(ids), 1)
            taus.append(float(np.mean(signs[iu])))
        out[name] = {
            "m": m,
            "memory_ratio_m_over_d": m / d,
            "recall@1": recall_at(got, truth, 1),
            "recall@10": recall_at(got, truth, 10),
            "top20_kendall": float(np.mean(taus)),
            "query_seconds": time.perf_counter() - t0,
        }
    return out


def main():
    verification = []
    for d, m in [(64, 4), (64, 16), (64, 32), (128, 8), (128, 32), (128, 64), (256, 32), (256, 64)]:
        empirical, agreement = pairwise_order_check(d, m, n_pairs=20000, seed=d + m)
        verification.append({
            "d": d, "m": m, "ratio": m / d,
            "empirical_tau": empirical,
            "empirical_agreement": agreement,
            "exact_tau": exact_tau(m, d),
            "asymptotic_tau": asymptotic_tau(m, d),
        })
    adaptive_m, adaptive_history = adaptive_dimension(128, target_tau=0.5)
    result = {"seed": SEED, "verification": verification,
              "adaptive_selection": {"final_m": adaptive_m, "history": adaptive_history},
              "retrieval": retrieval_trial()}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()