# ============================================================
# 00 — Hello, world
#
# Colab:  !pip install winnex-madhava
#
# WHAT THIS SHOWS
#   The smallest possible working example. Six lines of real work.
#   No database, no embeddings provider, no configuration to explain.
#
#   Run this first. If it prints 0 bound violations, the install is
#   good and every other example will work.
# ============================================================

# 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 numpy as np
import winnex_madhava as mad

# 1. a corpus of 5 vectors with 16 dimensions
corpus = np.array([
    [1.0, 0.0, 0.0, 0.0] + [0.0] * 12,    # document 0
    [0.0, 1.0, 0.0, 0.0] + [0.0] * 12,    # document 1
    [0.0, 0.0, 1.0, 0.0] + [0.0] * 12,    # document 2
    [0.0, 0.0, 0.0, 1.0] + [0.0] * 12,    # document 3
    [0.7, 0.7, 0.0, 0.0] + [0.0] * 12,    # document 4
], dtype=np.float32)

# 2. build the index
engine = mad.build_engine(corpus, dim=16, k=2, k1_fraction=1.0,
                          postfilter=True, audit_exhaustive=True)

# 3. search with a query identical to document 0
query = corpus[0]
result = mad.retrieve(engine, query, k=2)

print("indices:", result["indices"])
print("bound violations:", result["bound_violations"])
print("latency:", f"{result['latency_ms']:.3f} ms")
print("proof:", result["proof"])

# the engine found document 0 first, and proved the other four
# could not be in the top-2.

assert result["indices"][0] == 0
assert result["bound_violations"] == 0

print("\nOK — working install.")