# !pip install -q pyarrow numpy sentence-transformers datasets sqlite-vec tqdm psutil

import gc
import hashlib
import json
import os
import pickle
import random
import re
import sqlite3
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path
from typing import Optional

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import psutil
import sqlite_vec
from sentence_transformers import SentenceTransformer
from tqdm import tqdm

random.seed(42)
np.random.seed(42)

EMBEDDING_MODEL = "all-MiniLM-L6-v2"  # 384-dim, fast, good enough
EMBEDDING_DIM = 384
INDEX_DIR = Path("./nanomem_index")
INDEX_DIR.mkdir(exist_ok=True)
CACHE_DIR = Path("./bench_cache")
CACHE_DIR.mkdir(exist_ok=True)

PROC = psutil.Process(os.getpid())


def mem_mb() -> float:
    """Current RSS in MB. Useful for the writeup section."""
    return PROC.memory_info().rss / 1024 / 1024


class MemoryBudget:
    """Track peak memory (RSS) across a block of code plus across the whole run.

    Use as a context manager around any code where memory matters.
    Required at Tier 93 (target 2048) and Tier 100 (target 1024).

        with MemoryBudget("Phase 1 build", limit_mb=2048) as mb:
            inv_idx.build(docs)
            mb.update()                # sample again after allocation-heavy steps
        # prints: [mem] Phase 1 build: peak 1340 MB / budget 2048 MB

    At the end of the run:
        print(f"Run peak: {MemoryBudget.run_peak():.0f} MB")

    The eval harness uses this automatically and reports peak_mb in
    every result dict, so your scoreboard table can quote it.

    Note: this reports `psutil` peak RSS, not the conceptual working set.
    At Tier 93/100, peak RSS will often read higher than 2 GB even when
    you only have one shard live at a time, because Python doesn't return
    freed pages to the OS. We grade single-shard residency by reading
    your notebook; this number is a soft signal. **Keep the `[mem]`
    prints visible in the submitted notebook.**
    """

    _run_peak: float = 0.0

    def __init__(self, label: str, limit_mb: float | None = None):
        self.label = label
        self.limit_mb = limit_mb
        self.start: float = 0.0
        self.peak: float = 0.0

    def __enter__(self):
        self.start = mem_mb()
        self.peak = self.start
        return self

    def update(self) -> float:
        cur = mem_mb()
        self.peak = max(self.peak, cur)
        MemoryBudget._run_peak = max(MemoryBudget._run_peak, self.peak)
        return cur

    def __exit__(self, *args):
        self.update()
        msg = f"[mem] {self.label}: peak {self.peak:.0f} MB"
        if self.limit_mb is not None:
            msg += f" / budget {self.limit_mb:.0f} MB"
            if self.peak > self.limit_mb:
                msg += " OVER BUDGET"
        print(msg)
        return False  # do not suppress exceptions

    @classmethod
    def run_peak(cls) -> float:
        """Highest RSS seen inside any MemoryBudget block so far."""
        return cls._run_peak

    @classmethod
    def reset(cls):
        cls._run_peak = 0.0


def memory_breakdown(components: dict, label: str = "memory profile") -> dict:
    """Print a per-component memory breakdown table.

    Use this to figure out *where* your memory is going before you decide
    what to optimize at Tier 93 or Tier 100. Pass a dict of
    {label: object}; values can be:

      - numpy.ndarray:    uses .nbytes (exact)
      - InvertedIndex:    walks postings dict + doc_lengths
      - LSHIndex:         walks hyperplane arrays + bucket dicts
      - int / float:      treated as MB directly (use for measurements
                          you took elsewhere, e.g., a SQLite file size)
      - anything else:    sys.getsizeof(obj) / 1024 / 1024 (approximate)

    Returns the {label: mb} dict so you can also persist it.

    **Keep the printed table visible in your notebook for Tier 93/100.**
    It is the artifact we read to verify residency.

    **Approximate numbers at Tier 87** (1 GB, 250K docs per dataset):

        COMPONENT                   MB     %
        embedding matrix          380    27
        inverted postings         400    28
        LSH bands                 300    21
        doc texts                 250    18
        vocabulary keys            30     2
        -----------------------------------
        TOTAL                   ~1400 MB

    Each row scales linearly with dataset size: ~5x at Tier 93, ~80x at
    Tier 100. That is why HashPartition is required at Tier 93+ -- you
    can only afford to hold ONE shard's slice of each row at a time.

    Tier targets after sharding (single-shard residency):

        Tier        Working set     Target     What to do
        Tier 87       ~1.4 GB     unlimited    no sharding needed
        Tier 93       ~1.7 GB        2 GB      N=4 shards is plenty
        Tier 100      ~700 MB        1 GB      N=128 + streaming
                                               embed + bounded vocab
    """
    import sys
    sizes = {}
    for k, v in components.items():
        if isinstance(v, np.ndarray):
            sizes[k] = v.nbytes / 1024 / 1024
        elif isinstance(v, (int, float)):
            sizes[k] = float(v)
        elif hasattr(v, 'postings') and hasattr(v, 'doc_lengths'):
            n = sys.getsizeof(v.postings)
            for term, plist in v.postings.items():
                n += sys.getsizeof(term) + sys.getsizeof(plist) + len(plist) * 56
            n += sys.getsizeof(v.doc_lengths)
            sizes[k] = n / 1024 / 1024
        elif hasattr(v, 'bands'):
            n = 0
            for hp, buckets in v.bands:
                n += hp.nbytes if isinstance(hp, np.ndarray) else 0
                n += sys.getsizeof(buckets)
                for _bid, doc_ids in buckets.items():
                    n += sys.getsizeof(doc_ids) + len(doc_ids) * 28
            sizes[k] = n / 1024 / 1024
        else:
            sizes[k] = sys.getsizeof(v) / 1024 / 1024

    total = sum(sizes.values())
    print(f"\n=== {label} ===")
    print(f"{'COMPONENT':<28} {'MB':>10} {'%':>6}")
    print("-" * 46)
    for k in sorted(sizes, key=lambda x: -sizes[x]):
        mb = sizes[k]
        pct = (100 * mb / total) if total else 0
        print(f"{k:<28} {mb:>10.1f} {pct:>5.1f}%")
    print("-" * 46)
    print(f"{'TOTAL':<28} {total:>10.1f}  100.0%")
    print(f"process peak so far:        {MemoryBudget.run_peak():>7.1f} MB")
    return sizes

_embed_model = None


def get_embedder() -> SentenceTransformer:
    global _embed_model
    if _embed_model is None:
        print("Loading embedding model (one-time)...")
        _embed_model = SentenceTransformer(EMBEDDING_MODEL)
    return _embed_model


def embed_text(text: str) -> np.ndarray:
    """Return a (EMBEDDING_DIM,) float32 unit vector."""
    v = get_embedder().encode([text], convert_to_numpy=True)[0].astype(np.float32)
    return v / (np.linalg.norm(v) + 1e-12)


def embed_batch(texts: list[str], batch_size: int = 128) -> np.ndarray:
    """Return an (N, EMBEDDING_DIM) float32 matrix of unit vectors."""
    M = get_embedder().encode(
        texts, batch_size=batch_size, show_progress_bar=False, convert_to_numpy=True
    ).astype(np.float32)
    norms = np.linalg.norm(M, axis=1, keepdims=True) + 1e-12
    return M / norms


def cached_embed_batch(doc_ids: list[str], texts: list[str], cache_path: Path) -> np.ndarray:
    """Compute embeddings once, save to Parquet, reload on subsequent calls.

    Use this everywhere instead of raw `embed_batch`. Embedding 250K documents
    takes ~7 minutes on Mac CPU; you do not want to redo it on every iteration.
    The cache key is `cache_path`; pick a name that includes the dataset and
    tier (e.g., `bench_cache/codesearchnet_1gb_embs.parquet`).

    On Colab specifically: Colab's local disk does NOT persist across session
    restarts. Mount Google Drive and point `cache_path` there:

        from google.colab import drive
        drive.mount('/content/drive')
        CACHE_DIR = Path("/content/drive/MyDrive/nanomem_cache")
        CACHE_DIR.mkdir(exist_ok=True)
        embs = cached_embed_batch(doc_ids, texts, CACHE_DIR / "codesearchnet_1gb.parquet")

    For the Wikipedia bonus tier where embedding takes hours, you may want to
    run the embedding pass on a separate machine with a beefier GPU and upload
    the resulting Parquet to Drive once.

    The cache uses a sidecar `<path>.key` file holding a sha1 fingerprint of
    your `doc_ids`. If the Parquet exists and the key matches, we reload (a
    few seconds). If the key is missing or stale, we re-embed and rewrite.
    This means changing the dataset auto-invalidates without you having to
    delete the cache by hand. The HIT / MISS line printed below tells you
    whether you actually paid the embedding cost on this run.
    """
    key = hashlib.sha1()
    key.update(str(len(doc_ids)).encode()); key.update(b"\x00")
    for did in doc_ids:
        key.update(did.encode()); key.update(b"\x00")
    key = key.hexdigest()
    key_path = cache_path.with_suffix(cache_path.suffix + ".key")
    if cache_path.exists() and key_path.exists() and key_path.read_text().strip() == key:
        _, cached_embs = load_embeddings_parquet(cache_path)
        print(f"  [embed cache HIT] {cache_path.name} ({cached_embs.shape[0]} docs)")
        return cached_embs
    reason = "no cache" if not cache_path.exists() else "doc_ids changed"
    print(f"  [embed cache MISS: {reason}] computing embeddings for {len(texts)} docs...")
    embs = embed_batch(texts)
    save_embeddings_parquet(doc_ids, embs, cache_path)
    key_path.write_text(key)
    print(f"  [embed cache WROTE] {cache_path.name}")
    return embs

_TOKEN_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")


def tokenize(text: str) -> list[str]:
    """Lowercase identifier-style tokenizer.

    Splits on non-identifier characters; keeps `verify_jwt` as one token.
    Modify if you want, but you do not have to.
    """
    return [t.lower() for t in _TOKEN_RE.findall(text)]

def save_documents_json(docs: list[dict], path: Path) -> None:
    with open(path, "w") as f:
        for d in docs:
            f.write(json.dumps(d) + "\n")


def load_documents_json(path: Path) -> list[dict]:
    docs = []
    with open(path) as f:
        for line in f:
            docs.append(json.loads(line))
    return docs


def save_embeddings_parquet(doc_ids: list[str], embeddings: np.ndarray, path: Path) -> None:
    """One Parquet file. doc_id (string) + embedding (fixed-size list of float32)."""
    table = pa.table({
        "doc_id": pa.array(doc_ids, type=pa.string()),
        "embedding": pa.array(
            [list(map(float, v)) for v in embeddings],
            type=pa.list_(pa.float32(), EMBEDDING_DIM),
        ),
    })
    pq.write_table(table, path, compression="snappy")


def load_embeddings_parquet(path: Path) -> tuple[list[str], np.ndarray]:
    table = pq.read_table(path)
    doc_ids = table.column("doc_id").to_pylist()
    embs = np.array(table.column("embedding").to_pylist(), dtype=np.float32)
    return doc_ids, embs


def brute_force_cosine(query_vec: np.ndarray, embeddings: np.ndarray, top_k: int = 10) -> list[int]:
    """Return the top_k indices into `embeddings` by cosine similarity."""
    scores = embeddings @ query_vec
    return list(np.argsort(-scores)[:top_k])

def random_projection_matrix(n_hyperplanes: int, dim: int = EMBEDDING_DIM) -> np.ndarray:
    """Returns (n_hyperplanes, dim) standard-normal hyperplanes."""
    return np.random.randn(n_hyperplanes, dim).astype(np.float32)


def hash_vector(vec: np.ndarray, hyperplanes: np.ndarray) -> int:
    """Project vec onto each hyperplane; return bit signature as int."""
    bits = (hyperplanes @ vec > 0).astype(np.uint8)
    sig = 0
    for b in bits:
        sig = (sig << 1) | int(b)
    return sig


def hash_batch(M: np.ndarray, hyperplanes: np.ndarray) -> np.ndarray:
    """Vectorized hash of an (N, dim) matrix. Returns (N,) int64 sigs."""
    bits = (M @ hyperplanes.T > 0).astype(np.uint8)
    n_hp = bits.shape[1]
    weights = (1 << np.arange(n_hp - 1, -1, -1)).astype(np.int64)
    return (bits.astype(np.int64) * weights).sum(axis=1)

def open_sqlite_db(path: Path) -> sqlite3.Connection:
    """Open a SQLite connection with the sqlite-vec extension loaded."""
    conn = sqlite3.connect(str(path))
    conn.enable_load_extension(True)
    sqlite_vec.load(conn)
    return conn


def build_sqlite_corpus(conn: sqlite3.Connection, docs: list[dict], embeddings: np.ndarray) -> None:
    """Set up FTS5 (BM25) + sqlite-vec (vector KNN) on a dataset.

    Creates three tables:
      docs_fts:   FTS5 full-text index for BM25
      docs_vec:   sqlite-vec virtual table for vector KNN
      doc_id_map: integer doc_idx -> string doc_id (for joining results)
    """
    conn.execute("CREATE VIRTUAL TABLE IF NOT EXISTS docs_fts USING fts5(doc_id UNINDEXED, content)")
    conn.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS docs_vec USING vec0(doc_idx INTEGER PRIMARY KEY, embedding float[{EMBEDDING_DIM}])")
    conn.execute("CREATE TABLE IF NOT EXISTS doc_id_map(doc_idx INTEGER PRIMARY KEY, doc_id TEXT)")
    conn.execute("BEGIN")
    conn.executemany("INSERT INTO docs_fts(doc_id, content) VALUES (?, ?)",
                     [(d["doc_id"], d["text"]) for d in docs])
    conn.executemany("INSERT INTO docs_vec(doc_idx, embedding) VALUES (?, ?)",
                     [(i, embeddings[i].tobytes()) for i in range(len(docs))])
    conn.executemany("INSERT INTO doc_id_map(doc_idx, doc_id) VALUES (?, ?)",
                     [(i, docs[i]["doc_id"]) for i in range(len(docs))])
    conn.execute("COMMIT")


def fts5_query(conn: sqlite3.Connection, query: str, top_k: int = 10) -> list[tuple[str, float]]:
    """BM25 query via FTS5. Returns [(doc_id, score), ...] sorted by relevance."""
    tokens = re.findall(r"[A-Za-z0-9_]+", query)
    if not tokens:
        return []
    fts_query = " OR ".join(f'"{t}"' for t in tokens)
    try:
        rows = conn.execute(
            "SELECT doc_id, -bm25(docs_fts) FROM docs_fts WHERE docs_fts MATCH ? ORDER BY rank LIMIT ?",
            (fts_query, top_k)).fetchall()
        return [(r[0], r[1]) for r in rows]
    except sqlite3.OperationalError:
        return []


def vec_query(conn: sqlite3.Connection, query_vec: np.ndarray, top_k: int = 10) -> list[tuple[str, float]]:
    """Brute-force vector KNN via sqlite-vec. Returns [(doc_id, similarity), ...]."""
    rows = conn.execute(
        """SELECT m.doc_id, -v.distance FROM docs_vec v
           JOIN doc_id_map m ON v.doc_idx = m.doc_idx
           WHERE v.embedding MATCH ? AND k = ? ORDER BY distance LIMIT ?""",
        (query_vec.tobytes(), top_k, top_k)).fetchall()
    return [(r[0], r[1]) for r in rows]

TIER_SIZES = {"tiny": 1_000, "100mb": 25_000, "1gb": 250_000, "5gb": 1_250_000, "20gb": 5_000_000}


def load_codesearchnet(tier: str = "tiny") -> tuple[list[dict], list[dict]]:
    """Load CodeSearchNet at the chosen tier. Returns (docs, eval_set).

    eval_set is built from docstrings: query is the first sentence of the
    docstring; relevant_doc_ids is [the function whose docstring it was].
    """
    from datasets import load_dataset

    if tier not in TIER_SIZES:
        raise ValueError(f"unknown tier {tier!r}; choose from {list(TIER_SIZES)}")
    n = TIER_SIZES[tier]
    use_streaming = (tier == "tiny")

    ds = load_dataset(
        "code-search-net/code_search_net", "python", split="train",
        streaming=use_streaming, trust_remote_code=True,
    )
    docs, eval_pool = [], []
    for row in ds:
        if len(docs) >= n:
            break
        code = row["func_code_string"]
        if not code:
            continue
        doc_id = f"python/{row['repository_name']}/{row['func_name']}::{len(docs)}"
        docs.append({"doc_id": doc_id, "text": code})
        docstring = (row.get("func_documentation_string") or "").strip()
        if docstring:
            first = next((line.strip() for line in docstring.split("\n") if line.strip()), "")
            if first and 20 <= len(first) <= 200:
                eval_pool.append({"query": first, "relevant_doc_ids": [doc_id]})
    rng = random.Random(42)
    eval_n = min(500, len(eval_pool))
    eval_set = rng.sample(eval_pool, eval_n)
    return docs, eval_set


def load_msmarco(tier: str = "tiny") -> tuple[list[dict], list[dict]]:
    """Load MS MARCO at the chosen tier. Returns (docs, eval_set).

    eval_set is real Bing/SQ queries from the BeIR/msmarco-qrels training split,
    filtered to those whose relevant docs are in the loaded subset.
    """
    from datasets import load_dataset

    if tier not in TIER_SIZES:
        raise ValueError(f"unknown tier {tier!r}; choose from {list(TIER_SIZES)}")
    n = TIER_SIZES[tier]

    corpus_ds = load_dataset("BeIR/msmarco", "corpus", split="corpus")
    queries_ds = load_dataset("BeIR/msmarco", "queries", split="queries")
    qrels_ds = load_dataset("BeIR/msmarco-qrels", split="train")

    docs = []
    docs_in_subset = set()
    for i, row in enumerate(corpus_ds):
        if i >= n:
            break
        doc_id = str(row["_id"])
        text = ((row.get("title") or "") + " " + (row.get("text") or "")).strip()
        docs.append({"doc_id": doc_id, "text": text})
        docs_in_subset.add(doc_id)

    qid_to_text = {str(row["_id"]): row["text"] for row in queries_ds}
    qid_to_relevant = {}
    for row in qrels_ds:
        qid = str(row["query-id"])
        did = str(row["corpus-id"])
        if did in docs_in_subset and int(row["score"]) > 0:
            qid_to_relevant.setdefault(qid, []).append(did)
    eligible = [q for q, r in qid_to_relevant.items() if r and q in qid_to_text]
    rng = random.Random(42)
    eval_n = min(500, len(eligible))
    eval_set = [{"query": qid_to_text[q], "relevant_doc_ids": qid_to_relevant[q]}
                for q in rng.sample(eligible, eval_n)]
    return docs, eval_set

def recall_at_k(retrieved: list[str], relevant: list[str], k: int = 5) -> float:
    return (1.0 if any(r in set(retrieved[:k]) for r in relevant) else 0.0) if relevant else 0.0


def mean_reciprocal_rank(retrieved: list[str], relevant: list[str]) -> float:
    rel_set = set(relevant)
    for i, doc in enumerate(retrieved, start=1):
        if doc in rel_set:
            return 1.0 / i
    return 0.0


def run_eval(query_fn, eval_set: list[dict], top_k: int = 10,
             memory_limit_mb: float | None = None) -> dict:
    """Run query_fn against an eval set. Reports R@1, R@5, R@10, MRR,
    p50/p95 latency, AND peak RSS during the run.

    Pass `memory_limit_mb=2048` for Tier 93 or `memory_limit_mb=1024`
    for Tier 100; the eval will assert your peak stayed under the
    cap. Omit to just observe peak without enforcement.
    """
    r1, r5, r10, mrrs, lat = [], [], [], [], []
    with MemoryBudget("eval run", limit_mb=memory_limit_mb) as mb:
        for ex in eval_set:
            t0 = time.perf_counter()
            results = query_fn(ex["query"])
            lat.append((time.perf_counter() - t0) * 1000)
            mb.update()
            retrieved = [d for d, _ in results] if (results and isinstance(results[0], tuple)) else (results or [])
            r1.append(recall_at_k(retrieved, ex["relevant_doc_ids"], 1))
            r5.append(recall_at_k(retrieved, ex["relevant_doc_ids"], 5))
            r10.append(recall_at_k(retrieved, ex["relevant_doc_ids"], 10))
            mrrs.append(mean_reciprocal_rank(retrieved, ex["relevant_doc_ids"]))
    return {
        "n": len(eval_set),
        "recall@1": float(np.mean(r1)),
        "recall@5": float(np.mean(r5)),
        "recall@10": float(np.mean(r10)),
        "mrr": float(np.mean(mrrs)),
        "p50_ms": float(np.percentile(lat, 50)),
        "p95_ms": float(np.percentile(lat, 95)),
        "peak_mb": mb.peak,
    }

TINY_CORPUS = [
    {"doc_id": "auth/jwt.py:1", "text": "def verify_jwt(token, secret): return jwt.decode(token, secret, algorithms=['HS256'])"},
    {"doc_id": "auth/jwt.py:2", "text": "def issue_jwt(user_id, secret, ttl=3600): return jwt.encode({'sub': user_id, 'exp': time.time() + ttl}, secret)"},
    {"doc_id": "auth/scopes.py:1", "text": "def verify_jwt_with_scope(token, secret, required_scope): payload = verify_jwt(token, secret); assert required_scope in payload['scopes']"},
    {"doc_id": "config/parser.py:1", "text": "def parse_config(path): with open(path) as f: return json.load(f)"},
    {"doc_id": "config/parser.py:2", "text": "def merge_configs(base, override): merged = dict(base); merged.update(override); return merged"},
    {"doc_id": "utils/load_json.py:1", "text": "def load_json(path): with open(path, 'r') as fh: return json.load(fh)"},
    {"doc_id": "utils/load_json.py:2", "text": "def save_json(obj, path): with open(path, 'w') as fh: json.dump(obj, fh, indent=2)"},
    {"doc_id": "db/connection.py:1", "text": "def connect_postgres(dsn): conn = psycopg2.connect(dsn); conn.autocommit = False; return conn"},
    {"doc_id": "db/connection.py:2", "text": "def close_connection(conn): conn.commit(); conn.close()"},
    {"doc_id": "db/transactions.py:1", "text": "def with_transaction(conn, fn): try: result = fn(conn); conn.commit(); return result; except: conn.rollback(); raise"},
    {"doc_id": "search/tokenize.py:1", "text": "def tokenize(text): return [w.lower() for w in re.split(r'\\W+', text) if w]"},
    {"doc_id": "search/bm25.py:1", "text": "def bm25_score(tf, df, n_docs, doc_len, avg_len, k1=1.2, b=0.75): idf = math.log((n_docs - df + 0.5) / (df + 0.5))"},
    {"doc_id": "cache/redis_client.py:1", "text": "def cache_get(key): return redis.get(key)"},
    {"doc_id": "cache/redis_client.py:2", "text": "def cache_set(key, value, ttl=300): redis.setex(key, ttl, value)"},
    {"doc_id": "models/user.py:1", "text": "class User: def __init__(self, id, email): self.id = id; self.email = email"},
    {"doc_id": "models/user.py:2", "text": "def get_user_by_email(session, email): return session.query(User).filter_by(email=email).first()"},
    {"doc_id": "api/routes.py:1", "text": "@app.route('/login', methods=['POST']) def login(): user = authenticate(request.json); return issue_jwt(user.id, SECRET)"},
    {"doc_id": "api/routes.py:2", "text": "@app.route('/logout', methods=['POST']) def logout(): revoke_token(request.headers['Authorization'])"},
    {"doc_id": "tests/test_auth.py:1", "text": "def test_jwt_roundtrip(): token = issue_jwt('user123', 'secret'); assert verify_jwt(token, 'secret')['sub'] == 'user123'"},
    {"doc_id": "tests/test_config.py:1", "text": "def test_parse_config(): cfg = parse_config('test.json'); assert cfg['db'] == 'postgres'"},
]

def setup_phase0(tier: str = "100mb") -> dict:
    """YOU IMPLEMENT.

    For each corpus (codesearchnet, msmarco):
      1. Load the corpus and eval set with the supplied loaders.
      2. Compute embeddings with embed_batch.
      3. Open a SQLite database at INDEX_DIR / f"{corpus}_phase0.db".
      4. Call build_sqlite_corpus(conn, docs, embs) to populate FTS5 + sqlite-vec.
      5. Run one test BM25 query and one test vector query; print the top result.

    Return a dict mapping corpus name to {"conn", "docs", "embs", "eval_set"}.
    """
    raise NotImplementedError("Phase 0: implement setup_phase0")


def run_phase_0_tests() -> None:
    state = setup_phase0(tier="100mb")
    assert "codesearchnet" in state and "msmarco" in state, "must set up both datasets"
    for corpus, s in state.items():
        assert s["conn"] is not None
        assert len(s["docs"]) > 1000, f"{corpus}: expected >1000 docs"
        assert s["embs"].shape[1] == EMBEDDING_DIM
        # FTS5 sanity check
        bm25_results = fts5_query(s["conn"], "function", top_k=3)
        assert len(bm25_results) > 0, f"{corpus}: FTS5 returned no results"
        # vec sanity check
        qv = embed_text("how do I authenticate users")
        vec_results = vec_query(s["conn"], qv, top_k=3)
        assert len(vec_results) > 0, f"{corpus}: sqlite-vec returned no results"
        print(f"  ✓ {corpus}: {len(s['docs']):,} docs loaded, FTS5+sqlite-vec working")
    print("\nPhase 0: ALL PASS")

class InvertedIndex:
    """YOU IMPLEMENT.

    Required attributes after build:
      self.postings:        dict[str, list[tuple[int, int]]]   # term -> [(doc_idx, tf), ...]
      self.doc_lengths:     list[int]                          # tokens per doc
      self.avg_doc_length:  float
      self.n_docs:          int
      self.doc_ids:         list[str]                          # doc_idx -> doc_id
    """

    def __init__(self) -> None:
        self.postings: dict[str, list[tuple[int, int]]] = {}
        self.doc_lengths: list[int] = []
        self.avg_doc_length: float = 0.0
        self.n_docs: int = 0
        self.doc_ids: list[str] = []

    def build(self, documents: list[dict]) -> None:
        """Ingest a list of {'doc_id': str, 'text': str} into the index."""
        # YOUR CODE HERE.
        # 1. Iterate documents, tokenize, count term frequencies.
        # 2. Append (doc_idx, tf) to self.postings[term] for each term.
        # 3. Track doc_lengths and doc_ids.
        # 4. Compute avg_doc_length.
        raise NotImplementedError("Phase 1: implement InvertedIndex.build")

    def df(self, term: str) -> int:
        return len(self.postings.get(term, []))

    def save(self, path: Path) -> None:
        with open(path, "wb") as f:
            pickle.dump({
                "postings": self.postings, "doc_lengths": self.doc_lengths,
                "avg_doc_length": self.avg_doc_length, "n_docs": self.n_docs,
                "doc_ids": self.doc_ids,
            }, f)

    @classmethod
    def load(cls, path: Path) -> "InvertedIndex":
        idx = cls()
        with open(path, "rb") as f:
            d = pickle.load(f)
        idx.postings = d["postings"]; idx.doc_lengths = d["doc_lengths"]
        idx.avg_doc_length = d["avg_doc_length"]; idx.n_docs = d["n_docs"]
        idx.doc_ids = d["doc_ids"]
        return idx


def bm25_query(index: InvertedIndex, query: str, top_k: int = 10,
               k1: float = 1.2, b: float = 0.75) -> list[tuple[str, float]]:
    """YOU IMPLEMENT.

    Tokenize the query, look up each term's posting list, score each doc that
    contains at least one query term using BM25, return top_k as
    [(doc_id, score), ...] sorted by score descending.
    """
    # YOUR CODE HERE.
    raise NotImplementedError("Phase 1: implement bm25_query")

def run_phase_1_tests() -> None:
    idx = InvertedIndex()
    idx.build(TINY_CORPUS)
    assert idx.n_docs == 20, f"expected 20 docs, got {idx.n_docs}"
    assert idx.df("verify_jwt") >= 2, f"expected df('verify_jwt')>=2, got {idx.df('verify_jwt')}"
    assert idx.df("redis") == 2, f"expected df('redis')==2, got {idx.df('redis')}"
    print("  ✓ inverted_index_build")

    save_path = INDEX_DIR / "phase1_test.pkl"
    idx.save(save_path)
    idx2 = InvertedIndex.load(save_path)
    assert idx2.n_docs == idx.n_docs
    assert idx2.df("verify_jwt") == idx.df("verify_jwt")
    print("  ✓ inverted_index_save_load")

    results = bm25_query(idx, "verify jwt token", top_k=5)
    assert len(results) > 0
    assert "auth/jwt.py:1" in [d for d, _ in results][:3], f"top-3 missing JWT doc: {results}"
    print("  ✓ bm25_basic_recall")

    results = bm25_query(idx, "verify_jwt_with_scope", top_k=5)
    assert results[0][0] == "auth/scopes.py:1", f"expected scopes doc first, got {results}"
    print("  ✓ bm25_identifier_match")

    results = bm25_query(idx, "kubernetes deployment yaml", top_k=5)
    assert results == [] or all(s == 0 for _, s in results)
    print("  ✓ bm25_no_match")

    print("\nPhase 1: ALL PASS")


def phase_1_parity(corpus: str = "codesearchnet", tier: str = "100mb") -> None:
    """Compare your BM25 against FTS5 on the same eval set.

    Reports the gap. Within ~5 percentage points of FTS5 Recall@5 is the bar.
    """
    print(f"  loading {corpus} {tier}...")
    if corpus == "codesearchnet":
        docs, eval_set = load_codesearchnet(tier)
    else:
        docs, eval_set = load_msmarco(tier)
    print(f"  {len(docs):,} docs, {len(eval_set):,} eval queries")

    inv_idx = InvertedIndex()
    inv_idx.build(docs)

    # SQLite reference
    db_path = INDEX_DIR / f"phase1_parity_{corpus}.db"
    if db_path.exists():
        db_path.unlink()
    conn = open_sqlite_db(db_path)
    embs_dummy = np.zeros((len(docs), EMBEDDING_DIM), dtype=np.float32)
    build_sqlite_corpus(conn, docs, embs_dummy)

    yours = run_eval(lambda q: bm25_query(inv_idx, q, top_k=10), eval_set)
    fts5 = run_eval(lambda q: fts5_query(conn, q, 10), eval_set)
    print(f"\n  Your BM25 R@5:  {yours['recall@5']:.3f}  (p50 {yours['p50_ms']:.1f} ms)")
    print(f"  FTS5 BM25 R@5:  {fts5['recall@5']:.3f}  (p50 {fts5['p50_ms']:.1f} ms)")
    gap = fts5['recall@5'] - yours['recall@5']
    print(f"  Gap:            {gap*100:+.1f} pts" + ("  ← within bar" if abs(gap) <= 0.05 else "  ← outside bar"))
    conn.close()

class LSHIndex:
    """YOU IMPLEMENT.

    Suggested structure:
      self.bands: list[tuple[np.ndarray, dict[int, list[int]]]]
        # for each band: (hyperplanes (bits_per_band, dim), signature -> [doc_idx, ...])
      self.embeddings: np.ndarray   # (N, dim) for exact re-rank
      self.doc_ids:    list[str]
    """

    def __init__(self, n_bands: int = 60, bits_per_band: int = 3) -> None:
        self.n_bands = n_bands
        self.bits_per_band = bits_per_band
        self.bands: list[tuple[np.ndarray, dict[int, list[int]]]] = []
        self.embeddings: np.ndarray = np.zeros((0, EMBEDDING_DIM), dtype=np.float32)
        self.doc_ids: list[str] = []

    def build(self, doc_ids: list[str], embeddings: np.ndarray) -> None:
        """YOUR CODE.

        1. Generate `n_bands` random hyperplane matrices (each (bits_per_band, dim)).
        2. Hash all embeddings against each band; bucket doc_idx by signature.
        3. Store self.embeddings, self.doc_ids, self.bands.
        """
        raise NotImplementedError("Phase 2: implement LSHIndex.build")

    def query(self, query_vec: np.ndarray, top_k: int = 10,
              candidate_budget: int = 1000) -> list[tuple[str, float]]:
        """YOUR CODE.

        1. Hash query_vec against each band; collect candidate doc_idx from any matching bucket.
        2. RANK candidates by hit count (how many bands they appeared in).
           Naive set-union-then-truncate breaks at scale. See FAQ.
        3. Take top `candidate_budget` candidates by hit count.
        4. Exact-cosine re-rank candidates against self.embeddings.
        5. Return top_k as [(doc_id, score), ...].
        """
        raise NotImplementedError("Phase 2: implement LSHIndex.query")

def run_phase_2_tests() -> None:
    docs = TINY_CORPUS
    embs = embed_batch([d["text"] for d in docs])

    idx = LSHIndex(n_bands=10, bits_per_band=4)
    idx.build([d["doc_id"] for d in docs], embs)
    assert len(idx.bands) == 10
    assert idx.embeddings.shape == (20, EMBEDDING_DIM)
    print("  ✓ lsh_build_no_crash")

    idx2 = LSHIndex(n_bands=20, bits_per_band=4)
    idx2.build([d["doc_id"] for d in docs], embs)
    target_idx = 0
    results = idx2.query(embs[target_idx], top_k=5)
    result_ids = [d for d, _ in results]
    assert docs[target_idx]["doc_id"] in result_ids, (
        f"self-retrieval failed: {docs[target_idx]['doc_id']} not in {result_ids}"
    )
    print("  ✓ lsh_self_retrieval")

    idx3 = LSHIndex(n_bands=30, bits_per_band=3)
    idx3.build([d["doc_id"] for d in docs], embs)
    queries = ["how do I parse a config file", "verify a JWT token", "redis cache"]
    overlaps = []
    for q in queries:
        qv = embed_text(q)
        bf_top = set(docs[i]["doc_id"] for i in brute_force_cosine(qv, embs, top_k=5))
        lsh_top = set(d for d, _ in idx3.query(qv, top_k=5))
        overlaps.append(len(bf_top & lsh_top) / 5)
    avg = sum(overlaps) / len(overlaps)
    assert avg >= 0.4, f"avg LSH/brute overlap {avg:.2f} too low"
    print(f"  ✓ lsh_recall_vs_brute_force (avg overlap {avg:.2f})")

    print("\nPhase 2: ALL PASS")


def phase_2_parity(corpus: str = "codesearchnet", tier: str = "100mb") -> None:
    """Compare your LSH against sqlite-vec's exact KNN on the same eval set."""
    print(f"  loading {corpus} {tier}...")
    if corpus == "codesearchnet":
        docs, eval_set = load_codesearchnet(tier)
    else:
        docs, eval_set = load_msmarco(tier)
    print(f"  {len(docs):,} docs, {len(eval_set):,} queries")
    embs = embed_batch([d["text"] for d in docs])

    lsh_idx = LSHIndex(n_bands=60, bits_per_band=3)
    lsh_idx.build([d["doc_id"] for d in docs], embs)

    db_path = INDEX_DIR / f"phase2_parity_{corpus}.db"
    if db_path.exists():
        db_path.unlink()
    conn = open_sqlite_db(db_path)
    build_sqlite_corpus(conn, docs, embs)

    yours = run_eval(lambda q: lsh_idx.query(embed_text(q), top_k=10), eval_set)
    vec = run_eval(lambda q: vec_query(conn, embed_text(q), 10), eval_set)
    print(f"\n  Your LSH R@5:        {yours['recall@5']:.3f}  (p50 {yours['p50_ms']:.1f} ms)")
    print(f"  sqlite-vec KNN R@5:  {vec['recall@5']:.3f}  (p50 {vec['p50_ms']:.1f} ms)")
    gap = vec['recall@5'] - yours['recall@5']
    print(f"  Gap:                 {gap*100:+.1f} pts" + ("  ← within bar" if abs(gap) <= 0.10 else "  ← outside bar"))
    conn.close()

def reciprocal_rank_fusion(ranked_lists: list[list[tuple[str, float]]],
                           k: int = 60, top_k: int = 10) -> list[tuple[str, float]]:
    """YOU IMPLEMENT.

    Standard RRF: for each (doc_id, score) appearing in any ranked_list at rank r,
    add 1.0 / (k + r) to its fused score. Sort, return top_k.
    """
    raise NotImplementedError("Phase 3: implement reciprocal_rank_fusion")


def hybrid_query(inv_idx: InvertedIndex, lsh_idx: LSHIndex, query: str,
                 top_k: int = 10) -> list[tuple[str, float]]:
    """YOU IMPLEMENT.

    Run BM25 + LSH, fuse with RRF, return top_k.
    """
    raise NotImplementedError("Phase 3: implement hybrid_query")


def query_planner(query: str) -> str:
    """YOU IMPLEMENT.

    Inspect the query text. Return one of: 'bm25', 'lsh', 'hybrid'.

    Hints:
    - Contains snake_case or camelCase identifiers? -> probably bm25
    - Looks like natural language ("how do I", "what is")? -> probably lsh
    - Mixed? -> hybrid
    """
    raise NotImplementedError("Phase 3: implement query_planner")


def custom_fusion(inv_idx: InvertedIndex, lsh_idx: LSHIndex, query: str,
                  top_k: int = 10) -> list[tuple[str, float]]:
    """YOU IMPLEMENT.

    A fusion variant of YOUR design that is meaningfully different from RRF.
    Examples: weighted score fusion (normalize then combine), confidence
    fall-through (use BM25 if it is confident, else fuse), agreement boost
    (RRF + extra weight when both signals agree on the top-1).

    Document what you did and why in a comment at the top of this function.
    """
    raise NotImplementedError("Phase 3: implement custom_fusion")

def run_phase_3_tests() -> None:
    a = [("doc1", 0.9), ("doc2", 0.5), ("doc3", 0.1)]
    b = [("doc2", 0.8), ("doc4", 0.7), ("doc1", 0.3)]
    fused = reciprocal_rank_fusion([a, b], k=60, top_k=4)
    fused_ids = [d for d, _ in fused]
    assert set(fused_ids[:2]) == {"doc1", "doc2"}, f"top-2 should be doc1,doc2: {fused_ids}"
    print("  ✓ rrf_basic")

    a2 = [("doc1", 0.9), ("doc2", 0.5)]
    fused2 = reciprocal_rank_fusion([a2], k=60, top_k=2)
    assert [d for d, _ in fused2] == ["doc1", "doc2"]
    print("  ✓ rrf_single_list")

    assert query_planner("verify_jwt_with_scope") == "bm25"
    assert query_planner("how do I parse a config file") == "lsh"
    print("  ✓ planner_routes")

    print("\nPhase 3: ALL PASS")

def run_two_corpus_benchmark(tier: str = "100mb") -> dict:
    """YOU IMPLEMENT.

    For each corpus in [codesearchnet, msmarco]:
      1. Load with the supplied loader.
      2. Compute embeddings.
      3. Build both indexes (InvertedIndex, LSHIndex).
      4. Run each retriever (bm25, lsh, hybrid, custom_fusion, planner-routed)
         through run_eval against the eval_set.
      5. Print a side-by-side table.

    Return a nested dict: {corpus_name: {retriever_name: metrics}}.
    """
    raise NotImplementedError("Phase 4: implement run_two_corpus_benchmark")


def run_phase_4_tests() -> None:
    """Sanity check the driver runs end-to-end on tiny."""
    out = run_two_corpus_benchmark(tier="tiny")
    assert "codesearchnet" in out and "msmarco" in out
    for corpus, retrievers in out.items():
        assert "bm25" in retrievers
        assert "lsh" in retrievers
        for name, m in retrievers.items():
            assert "recall@5" in m and "p50_ms" in m
    print("\nPhase 4: ALL PASS")

def run_production_scoreboard(tier: str = "100mb") -> dict:
    """YOU IMPLEMENT.

    For each corpus, on top of your Phase 4 results, add:
      - FTS5 BM25 metrics (use the supplied fts5_query)
      - sqlite-vec exact KNN metrics (use the supplied vec_query)

    Print a 7-column table with all retrievers side-by-side.

    Return a dict of the same shape as Phase 4, with the production rows added.
    """
    raise NotImplementedError("Phase 5: implement run_production_scoreboard")


def run_phase_5_tests() -> None:
    out = run_production_scoreboard(tier="tiny")
    for corpus, retrievers in out.items():
        assert "FTS5 BM25" in retrievers or "fts5_bm25" in retrievers, f"{corpus}: missing FTS5 row"
    print("\nPhase 5: ALL PASS")

def run_all_tests() -> None:
    print("=" * 60)
    print("RUNNING ALL PHASE TESTS")
    print("=" * 60)
    print("\nPhase 0:")
    run_phase_0_tests()
    print("\nPhase 1:")
    run_phase_1_tests()
    print("\nPhase 2:")
    run_phase_2_tests()
    print("\nPhase 3:")
    run_phase_3_tests()
    print("\nPhase 4:")
    run_phase_4_tests()
    print("\nPhase 5:")
    run_phase_5_tests()
    print("\n" + "=" * 60)
    print("ALL PHASES PASS")
    print("=" * 60)

CLI_TEMPLATE = '''
"""nanomem CLI. Usage:
    python nanomem.py index --corpus path/to/corpus/
    python nanomem.py query "how do I parse a config file"
"""
import argparse, json
from pathlib import Path

INDEX_DIR = Path("./nanomem_index")

def cmd_index(args):
    # Load corpus, build inverted index + LSH index, save to INDEX_DIR.
    raise NotImplementedError

def cmd_query(args):
    # Load indexes, run hybrid_query, print top-K as JSON.
    raise NotImplementedError

if __name__ == "__main__":
    p = argparse.ArgumentParser(prog="nanomem")
    sub = p.add_subparsers(dest="cmd", required=True)
    p_index = sub.add_parser("index"); p_index.add_argument("--corpus", required=True)
    p_query = sub.add_parser("query"); p_query.add_argument("query_text", type=str)
    args = p.parse_args()
    if args.cmd == "index": cmd_index(args)
    elif args.cmd == "query": cmd_query(args)
'''


if __name__ == "__main__":
    # Uncomment as you finish each phase:
    # run_phase_0_tests()
    # run_phase_1_tests()
    # run_phase_2_tests()
    # run_phase_3_tests()
    # run_phase_4_tests()
    # run_phase_5_tests()
    # run_all_tests()
    pass