# ============================================================
# 01 — Law firm: e-discovery with provable completeness
#
# Colab:  !pip install winnex-madhava
#
# THE REAL PROBLEM
#   A law firm receives a document production request. Counsel must
#   certify that the production is complete — and certify it under
#   penalty of sanctions for spoliation.
#
#   Every e-discovery platform uses approximate vector search. When
#   counsel is asked "did you search for X?", the answer is "yes"
#   but nobody can prove the search did not silently miss the
#   responsive document.
#
#   This is the scenario where Madhava changes what is possible: the
#   firm can produce the responsive set AND a machine-checkable
#   record that every excluded document was proven irrelevant.
#
# WHAT THIS EXAMPLE MODELS
#   A litigation support database (Relativity, Everlaw, Concordance
#   all store custodial documents in SQL). We load custodial emails,
#   index them, run a relevance query, and produce the audit record.
# ============================================================

# 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 = 256

def embed(text: str) -> np.ndarray:
    """Demo embedding — swap for a legal-domain encoder in production."""
    v = np.zeros(DIM, dtype=np.float32)
    for token in text.lower().split():
        for i, ch in enumerate(token):
            v[(ord(ch) * 31 + i * 7) % DIM] += 1.0
    n = np.linalg.norm(v)
    return v / n if n else v

# ── 1. custodial documents as they sit in litigation support ─
conn = sqlite3.connect(":memory:")
conn.execute("""
    CREATE TABLE documents (
        doc_id       TEXT PRIMARY KEY,
        custodian    TEXT,
        doc_date     TEXT,
        file_type    TEXT,
        subject      TEXT,
        body         TEXT,
        privilege    INTEGER DEFAULT 0
    )
""")

registros = [
    ("DOC-0001", "j.silva@acme.com.br", "2024-03-12", "email",
     "Re: Contrato de Fornecimento - Cláusula 7",
     "Conforme discutido, a cláusula 7 do contrato prevê multa de 2% ao mês "
     "em caso de atraso na entrega. O cliente solicitou revisão do prazo.", 0),

    ("DOC-0002", "m.costa@acme.com.br", "2024-03-14", "email",
     "Re: Contrato de Fornecimento - prazo",
     "Não concordo com a redução de prazo. A produção precisa de 90 dias. "
     "Se o cliente exigir 60, haverá risco de atraso e multa.", 0),

    ("DOC-0003", "j.silva@acme.com.br", "2024-04-02", "email",
     "Ata de reunião - diretoria",
     "A diretoria aprovou o pagamento da multa reclamada pelo cliente, "
     "no valor de R$ 48.000,00, sem reconhecimento de culpa.", 0),

    ("DOC-0004", "financeiro@acme.com.br", "2024-04-05", "planilha",
     "Conciliação de multas Q1 2024",
     "Multa paga em 04/04/2024: R$ 48.000,00. Classificação contábil: "
     "indenização contratual. Aprovação: diretoria.", 0),

    ("DOC-0005", "j.silva@acme.com.br", "2024-02-20", "email",
     "Política de férias 2024",
     "O período de férias coletivas será de 23/12 a 06/01, conforme "
     "comunicado da área de RH.", 0),

    ("DOC-0006", "rh@acme.com.br", "2024-01-15", "contrato",
     "Contrato de Trabalho - confidencial",
     "Cláusulas de confidencialidade e não concorrência aplicáveis ao "
     "cargo de diretor comercial.", 1),   # privileged

    ("DOC-0007", "m.costa@acme.com.br", "2024-03-28", "email",
     "Nota fiscal 4471",
     "NF 4471 referente ao fornecimento de março, valor R$ 1.240.000,00, "
     "vencimento em 30 dias.", 0),
]

for r in registros:
    conn.execute("INSERT INTO documents VALUES (?,?,?,?,?,?,?)", r)
conn.commit()

# ── 2. filter out privileged material BEFORE indexing ───────
# Privileged documents must never enter a searchable index — that is
# the first duty in a production review.
rows = conn.execute("""
    SELECT doc_id, custodian, subject, body
    FROM documents
    WHERE privilege = 0
    ORDER BY doc_id
""").fetchall()

doc_ids = [r[0] for r in rows]
textos = [f"{r[2]}. {r[3]}" for r in rows]
corpus = np.stack([embed(t) for t in textos]).astype(np.float32)

print(f"documents in database : {len(registros)}")
print(f"privileged withheld   : {len(registros) - len(rows)}")
print(f"indexed for search    : {len(rows)}\n")

# ── 3. build the index with the audit path on ───────────────
engine = mad.build_engine(
    corpus,
    dim=DIM,
    metric="cosine",
    k=3,
    k1_fraction=1.0,
    postfilter=True,
    audit_exhaustive=True,
)

# ── 4. the request: find everything about the delay penalty ─
request = "multa por atraso na entrega conforme cláusula contratual"

resultado = mad.retrieve(engine, embed(request), corpus_text=textos, k=3)

print(f"request: {request}\n")
print("responsive documents:")
q = embed(request)
for pos, idx in enumerate(resultado["indices"], 1):
    score = float(np.dot(q, corpus[idx]))
    print(f"  {pos}. [{score:.3f}] {doc_ids[idx]}  {rows[idx][2]}")

# ── 5. the certification a court can check ──────────────────
produzidos = [doc_ids[i] for i in resultado["indices"]]
retidos = [doc_ids[i] for i in range(len(doc_ids)) if i not in resultado["indices"]]

print(f"\n─── production certificate ───")
print(f"  production set        : {produzidos}")
print(f"  withheld as irrelevant: {retidos}")
print(f"  bound violations      : {resultado['bound_violations']}  (must be 0)")
print(f"  latency               : {resultado['latency_ms']:.2f} ms")

assert resultado["bound_violations"] == 0

print("""
What this certificate asserts
─────────────────────────────
For each document in the withheld set, the engine computed a
Cauchy–Schwarz upper bound below the score of the third responsive
document. That means the withheld documents could not belong in the
responsive set for this query — and the arithmetic is reproducible
by an opposing expert.

An approximate index gives you the same three documents and no way
to make that statement.
""")

# ── 6. why this matters in a spoliation challenge ───────────
print("Why counsel cares")
print("─" * 60)
print("""In discovery, you are asked not only "did you search" but
"can you demonstrate the search was complete". A bounding proof is
the difference between asserting completeness and evidencing it.
""")