# ============================================================
# 03 — Reproduce the BIGANN benchmark yourself
#
# Colab:  !pip install winnex-madhava
#
# WHAT THIS SHOWS
#   The claim on the website is that Madhava reproduces the exact
#   scan while staying fast. This script proves the first half of
#   that claim on synthetic data you generate yourself, so you do
#   not have to trust our numbers or download a 100M dataset.
#
#   It measures:
#     · top-1 agreement with an exact cosine scan
#     · setR@10 overlap with the exact top-10
#     · bound violations (must be zero)
#
# WHY SYNTHETIC FIRST
#   If a search engine cannot reproduce exact scan on data you made,
#   it will not do so on BIGANN. Start here, then move to a real
#   dataset. The Kaggle notebooks use the real BIGANN ground truth.
# ============================================================

# instalação automática — Colab, Jupyter e script local.
# No Colab, "!pip" dentro de comentário não executa: a instalação
# precisa acontecer de dentro da célula, antes dos imports.
import importlib
import subprocess
import sys

for _pkg in ["winnex-madhava"]:
    _mod = _pkg.replace("-", "_")
    try:
        __import__(_mod)
        continue
    except ImportError:
        pass
    _r = subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q", "--upgrade", _pkg],
        capture_output=True, text=True
    )
    # O interpretador do Colab já está em execução quando o pip termina.
    # Sem invalidar o cache, o import abaixo pode não enxergar o pacote.
    importlib.invalidate_caches()
    try:
        __import__(_mod)
    except ImportError:
        sys.stderr.write(_r.stdout or "")
        sys.stderr.write(_r.stderr or "")
        raise SystemExit(
            "\n\nNão foi possível importar " + _pkg + " após instalar.\n"
            "No Colab: menu 'Ambiente de execução' -> 'Reiniciar sessão',\n"
            "e execute esta célula novamente.\n"
        )

import time
import numpy as np
import winnex_madhava as mad

N = 20_000        # corpus size
D = 384           # dimensions
K = 10            # neighbours to retrieve
QUERIES = 50

rng = np.random.default_rng(42)

print(f"generating {N:,} vectors of {D} dimensions ...")
corpus = rng.standard_normal((N, D), dtype=np.float32)
corpus /= np.linalg.norm(corpus, axis=1, keepdims=True)
queries = rng.standard_normal((QUERIES, D), dtype=np.float32)
queries /= np.linalg.norm(queries, axis=1, keepdims=True)

# ── exact ground truth ──────────────────────────────────────

# brute-force cosine scan. This is the ceiling, not an opinion.
print("computing exact ground truth (brute force) ...")
t0 = time.perf_counter()
exact_top = []
for q in queries:
    scores = corpus @ q
    exact_top.append(np.argsort(-scores)[:K])
exact_seconds = time.perf_counter() - t0

# ── build the engine ────────────────────────────────────────
print("building Madhava engine ...")
t0 = time.perf_counter()
engine = mad.build_engine(
    corpus,
    dim=D,
    metric="cosine",
    k=K,
    k1_fraction=1.0,
    postfilter=True,
    audit_exhaustive=True,
)
build_seconds = time.perf_counter() - t0

# ── query ───────────────────────────────────────────────────
print("searching ...")
t0 = time.perf_counter()
pred_top = []
violations = 0
for q in queries:
    r = mad.retrieve(engine, q, k=K)
    pred_top.append(np.asarray(r["indices"]))
    violations += int(r["bound_violations"])
search_seconds = time.perf_counter() - t0

# ── metrics ─────────────────────────────────────────────────
top1 = np.mean([pred_top[i][0] == exact_top[i][0] for i in range(QUERIES)])
setr = np.mean([
    len(set(pred_top[i]) & set(exact_top[i])) / K for i in range(QUERIES)
])

print("\n" + "=" * 58)
print(f"corpus            {N:,} x {D}")
print(f"queries           {QUERIES}")
print(f"k                 {K}")
print("-" * 58)
print(f"top-1 agreement   {top1:.4f}")
print(f"setR@{K} overlap   {setr:.4f}")
print(f"bound violations  {violations}   (must be 0)")
print("-" * 58)
print(f"build time        {build_seconds*1000:.1f} ms")
print(f"search time       {search_seconds*1000:.1f} ms "
      f"({search_seconds/QUERIES*1000:.2f} ms per query)")
print(f"brute force time  {exact_seconds*1000:.1f} ms")
print("=" * 58)

assert violations == 0, "bound violations detected — engine is unsound"

if top1 == 1.0 and setr == 1.0:
    print("\nOK — the engine reproduced the exact scan on every query.")
else:
    print(f"\nNote: agreement is {top1:.3f} / {setr:.3f}. With audit_exhaustive")
    print("and k1_fraction=1.0 it should be 1.0000. If it is not, increase")
    print("k2_max or open an issue with the configuration above.")

# ── why this matters ────────────────────────────────────────
print("""
What you just measured
----------------------
An approximate index (HNSW, IVF, PQ) would return a top-10 that
overlaps the exact answer by some percentage — and give you no way
to know which queries were wrong.

The engine here returned the exact answer and recorded, for every
document not returned, the bound proving it could not belong.
That record is what an auditor can check.
""")