"""Build a 10,000-document index and price the same boolean queries under
two intersection strategies.

The corpus is synthetic: each term is sprinkled into documents with a fixed
probability from a seeded RNG. That is not what English looks like in
detail, but it gives the one property this demo needs -- a vocabulary whose
document frequencies span four orders of magnitude, pinned to exact values
that a test can assert.
"""

import random
from functools import partial

from inverted_index import (
    build_index,
    document_frequency,
    intersect_gallop,
    intersect_linear,
    query,
)

# (term, probability that any given document contains it). Ordered from
# stopword-common to vanishingly rare; the RNG is consumed in exactly this
# order, which is what pins every document frequency below.
VOCAB = (
    ("the", 0.85),
    ("of", 0.60),
    ("and", 0.45),
    ("data", 0.30),
    ("ranking", 0.20),
    ("search", 0.12),
    ("index", 0.06),
    ("postings", 0.02),
    ("corpus", 0.008),
    ("relevance", 0.004),
    ("obelisk", 0.0003),
    ("quantum", 0.0002),
)

N_DOCS = 10_000
SEED = 42

# The headline counterfactual, made runnable: galloping with the swap turned
# off, so the loop is driven by whichever list the query happened to name
# first instead of by the shorter one.
GALLOP_FROM_LONGEST = partial(intersect_gallop, drive_shortest=False)


def make_corpus(n_docs=N_DOCS, seed=SEED):
    """One `random.Random(seed)`, one draw per (document, vocabulary term)."""
    rng = random.Random(seed)
    docs = []
    for _ in range(n_docs):
        words = []
        for term, p in VOCAB:
            r = rng.random()
            if r < p:
                words.append(term)
                if r < p * 0.3:
                    words.append(term)  # real text repeats terms; so does this
        docs.append(" ".join(words))
    return docs


QUERIES = (
    ("the",),
    ("the", "quantum"),
    ("the", "obelisk"),
    ("search", "index"),
    ("of", "the"),
)


def main():
    docs = make_corpus()
    index = build_index(docs)

    print(f"corpus: {len(docs)} documents, {len(index)} distinct terms\n")

    print("document frequencies (= postings list lengths)")
    print(f"  {'term':<12}{'df':>8}{'% of corpus':>14}")
    for term, _ in VOCAB:
        df = document_frequency(index, term)
        print(f"  {term:<12}{df:>8}{df / len(docs):>13.2%}")

    print()
    print("query cost, in postings entries examined")
    print(
        f"  {'query':<22}{'hits':>7}{'gallop':>9}{'linear':>9}"
        f"{'no-swap':>9}{'best/worst':>12}"
    )
    for terms in QUERIES:
        label = " AND ".join(terms)
        g_hits, g_steps = query(index, terms, intersect_gallop)
        l_hits, l_steps = query(index, terms, intersect_linear)
        n_hits, n_steps = query(index, terms, GALLOP_FROM_LONGEST)
        assert g_hits == l_hits == n_hits, (terms, g_hits, l_hits, n_hits)
        costs = (g_steps, l_steps, n_steps)
        spread = f"{max(costs) / min(costs):.0f}x"
        print(
            f"  {label:<22}{len(g_hits):>7}{g_steps:>9}{l_steps:>9}"
            f"{n_steps:>9}{spread:>12}"
        )

    print()
    print("AND-ing every term against 'the', from commonest to rarest")
    print(f"  {'term':<12}{'df':>7}{'skew':>10}{'gallop':>9}{'linear':>9}   winner")
    for term, _ in VOCAB[1:]:
        df = document_frequency(index, term)
        _, g_steps = query(index, ("the", term), intersect_gallop)
        _, l_steps = query(index, ("the", term), intersect_linear)
        skew = document_frequency(index, "the") / df
        verdict = "gallop" if g_steps < l_steps else "linear"
        print(
            f"  {term:<12}{df:>7}{skew:>9.0f}:1{g_steps:>9}{l_steps:>9}"
            f"   {verdict} wins"
        )

    print()
    print("the opener -- adding a term can make the query CHEAPER:")
    for terms in (("the",), ("the", "quantum"), ("the", "search")):
        hits, steps = query(index, terms, intersect_gallop)
        print(f"  {' AND '.join(terms):<22}{len(hits):>6} hits  {steps:>6} steps")
    _, solo = query(index, ("the",), intersect_gallop)
    _, pair = query(index, ("the", "quantum"), intersect_gallop)
    print(f"  ratio: {solo / pair:.0f}x fewer steps for the MORE constrained query")
    print("  ...but only if the added term is selective -- 'search' is not.")

    print()
    print("THE HEADLINE -- same query, same 5 hits, one line different:")
    g_hits, g_steps = query(index, ("the", "obelisk"), intersect_gallop)
    l_hits, l_steps = query(index, ("the", "obelisk"), intersect_linear)
    n_hits, n_steps = query(index, ("the", "obelisk"), GALLOP_FROM_LONGEST)
    print(f"  drive from the SHORT list {len(g_hits):>4} hits  {g_steps:>6} steps")
    print(f"  walk both lists (linear)  {len(l_hits):>4} hits  {l_steps:>6} steps")
    print(f"  drive from the LONG list  {len(n_hits):>4} hits  {n_steps:>6} steps")
    print(f"  short vs long: {n_steps / g_steps:.0f}x, and the only difference is")
    print("    `if drive_shortest and len(a) > len(b): a, b = b, a`")

    print()
    print("the boundary -- two common terms, and the bet inverts:")
    g_hits, g_steps = query(index, ("of", "the"), intersect_gallop)
    l_hits, l_steps = query(index, ("of", "the"), intersect_linear)
    print(f"  drive from the SHORT list {len(g_hits):>4} hits  {g_steps:>6} steps")
    print(f"  walk both lists (linear)  {len(l_hits):>4} hits  {l_steps:>6} steps")
    print(f"  ratio: {g_steps / l_steps:.1f}x WORSE for galloping")


if __name__ == "__main__":
    main()