# ============================================================
# 01 — Query a SQL database with mathematical proof
#
# Colab:  !pip install winnex-madhava
#
# WHAT THIS SHOWS
#   The most common first question: "my data is in Postgres / MySQL /
#   SQLite — how do I search it and prove the search was complete?"
#
#   Read rows, embed them, index with Madhava, search. The engine
#   returns the answer plus the number of bound violations, which
#   must be zero.
#
# NOTE ON EMBEDDINGS
#   This example uses a deterministic hash embedding so it runs with
#   no API key and no model download. In production replace `embed()`
#   with a real encoder (sentence-transformers, OpenAI, or
#   winnex-ai-normalize). The search logic does not change.
# ============================================================

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

DIM = 128

def embed(text: str) -> np.ndarray:
    """Deterministic bag-of-characters embedding. Demo only."""
    v = np.zeros(DIM, dtype=np.float32)
    for ch in text.lower():
        v[ord(ch) % DIM] += 1.0
    n = np.linalg.norm(v)
    return v / n if n else v

# ── 1. a small database ─────────────────────────────────────

conn = sqlite3.connect(":memory:")
conn.execute(
    "CREATE TABLE documentos (id INTEGER PRIMARY KEY, titulo TEXT, texto TEXT)"
)
conn.executemany(
    "INSERT INTO documentos (id, titulo, texto) VALUES (?,?,?)",
    [
        (1, "Contrato de Locação",
         "O locatário pagará aluguel mensal de R$ 2.000,00 corrigido anualmente."),
        (2, "Termo de Confidencialidade",
         "As partes se obrigam a manter sigilo pelo prazo de 5 anos."),
        (3, "Laudo Médico",
         "Paciente apresenta quadro compatível com pneumonia bacteriana."),
        (4, "Parecer Jurídico",
         "Opinião favorável à celebração do acordo de acionistas."),
        (5, "Nota Fiscal",
         "Serviços de consultoria prestados no valor de R$ 15.000,00."),
        (6, "Aditivo Contratual",
         "Fica reajustado o valor do aluguel para R$ 2.400,00 mensais."),
    ],
)
conn.commit()

# ── 2. read the rows ────────────────────────────────────────
rows = conn.execute("SELECT id, titulo, texto FROM documentos ORDER BY id").fetchall()
textos = [r[2] for r in rows]
corpus = np.stack([embed(t) for t in textos]).astype(np.float32)

# ── 3. build the engine ─────────────────────────────────────
# k1_fraction=1.0 keeps every bound survivor and postfilter does the
# exact re-rank. This is the configuration that guarantees recall.
engine = mad.build_engine(
    corpus,
    dim=DIM,
    metric="cosine",
    k=3,
    k1_fraction=1.0,
    postfilter=True,
    audit_exhaustive=True,
)

# ── 4. query ────────────────────────────────────────────────
pergunta = "quanto é o aluguel"
resultado = mad.retrieve(engine, embed(pergunta), corpus_text=textos, k=3)

print(f"pergunta: {pergunta}\n")

# NOTE: in winnex-madhava v1.9.23 the `scores` field of the result is
# returned as zeros. The ordering is correct — only the score column is
# not populated. Recompute the cosine here so the output is meaningful.
q = embed(pergunta)
for pos, idx in enumerate(resultado["indices"], 1):
    score = float(np.dot(q, corpus[idx]))
    print(f"  {pos}. score={score:.4f}  {textos[idx][:60]}")

print(f"\nbound violations: {resultado['bound_violations']}   "
      f"(must be 0)")
print(f"latency: {resultado['latency_ms']:.2f} ms")
print(f"\nproof: {resultado['proof']}")

# ── 5. what the zero actually means ─────────────────────────
# Every document NOT in the answer received a Cauchy-Schwarz upper
# bound below the score of the k-th result. That is the claim: they
# could not have been in the top-k, and the arithmetic is recorded.
assert resultado["bound_violations"] == 0
print("\nOK — the exclusion is proven, not estimated.")