"""A miniature inverted index: tokenize documents, build postings lists, and
answer boolean AND queries two different ways.

The index itself is the boring half. The interesting half is the
*intersection*: given two sorted lists of document ids, which one do you walk,
and how? Every function that touches a postings list is handed a `Steps`
meter, so the cost of a query is a number the demo can print rather than
something you have to profile for.

One step = one postings entry examined. That is the unit the whole toy is
denominated in.

Document ids are assigned in corpus order, so postings lists come out sorted
by construction -- there is no hashing and no sort anywhere in this file.
"""


class Steps:
    """A cost meter, passed down into every list walk.

    Kept as an explicit object rather than a global counter so two strategies
    can be measured against the same query in the same process without one
    contaminating the other.
    """

    def __init__(self):
        self.n = 0

    def tick(self, k=1):
        self.n += k


def tokenize(text):
    """Lowercase and split on whitespace. That is the entire analyzer.

    No stemming, no stopword list -- keeping "the" in the index is what makes
    the cost difference in the demo visible.
    """
    return text.lower().split()


def build_index(docs):
    """Map each term to the sorted list of document ids containing it.

    `docs` is a sequence; a document's id is its position in that sequence.
    Because ids are handed out in increasing order and each document is
    visited once, each postings list is built in ascending order and never
    needs sorting. The `postings[-1] != doc_id` guard collapses repeated
    occurrences of a term within one document into a single entry, which is
    what keeps the list strictly increasing.
    """
    index = {}
    for doc_id, text in enumerate(docs):
        for term in tokenize(text):
            postings = index.setdefault(term, [])
            if not postings or postings[-1] != doc_id:
                postings.append(doc_id)
    return index


def document_frequency(index, term):
    """How many documents contain `term`. Also the length of its postings."""
    return len(index.get(term, ()))


def intersect_linear(a, b, steps):
    """Classic sort-merge intersection: two pointers, always advance the one
    pointing at the smaller id.

    Symmetric, and completely indifferent to which list is shorter. Cost is
    bounded by the position at which the first list runs out -- which, when
    one list is tiny and the other is huge, is still most of the huge one.
    """
    out = []
    i = j = 0
    while i < len(a) and j < len(b):
        steps.tick()
        if a[i] == b[j]:
            out.append(a[i])
            i += 1
            j += 1
        elif a[i] < b[j]:
            i += 1
        else:
            j += 1
    return out


def gallop_index(postings, start, target, steps):
    """Return the smallest index >= `start` whose posting is >= `target`
    (or len(postings) if there is none), by exponential probe then binary
    search.

    The exponential phase is what makes this cheaper than a plain binary
    search over the whole list: it finds a bracket whose width is
    proportional to the *distance actually travelled*, not to the list's
    length. A target two entries away costs a couple of probes no matter how
    long the list is.
    """
    n = len(postings)
    if start >= n:
        return n
    steps.tick()
    if postings[start] >= target:
        return start

    bound = 1
    while start + bound < n:
        steps.tick()
        if postings[start + bound] >= target:
            break
        bound *= 2

    # Everything at or below start + bound//2 is known to be < target, and
    # the answer is at or below start + bound (or past the end of the list).
    lo = start + bound // 2 + 1
    hi = min(start + bound, n - 1)
    while lo <= hi:
        steps.tick()
        mid = (lo + hi) // 2
        if postings[mid] < target:
            lo = mid + 1
        else:
            hi = mid - 1
    return lo


def intersect_gallop(a, b, steps, drive_shortest=True):
    """Intersection that drives the loop from the *shorter* list, skipping
    through the longer one with `gallop_index`.

    The swap is the whole bet: the number of probes is set by the length of
    the driving list, so making that the short one turns a walk of the long
    list into a handful of jumps over it.

    `drive_shortest=False` turns the swap off. It exists only so the demo can
    price the identical query with and without that one line; a real engine
    has no reason ever to pass it.
    """
    if drive_shortest and len(a) > len(b):
        a, b = b, a

    out = []
    j = 0
    for doc in a:
        steps.tick()
        j = gallop_index(b, j, doc, steps)
        if j == len(b):
            break
        if b[j] == doc:
            out.append(doc)
            j += 1
    return out


def query(index, terms, intersect=intersect_gallop):
    """Answer `terms[0] AND terms[1] AND ...`, returning (doc_ids, steps).

    A single-term query has no intersection to do, so its cost is just the
    cost of handing back every posting -- one step each. A multi-term query
    never materializes its first list; it only ever touches the entries the
    intersection asks for. That asymmetry is the point: adding a term
    replaces "emit everything" with "probe a little".
    """
    steps = Steps()
    lists = [index.get(term, []) for term in terms]
    if not lists:
        return [], steps.n
    if len(lists) == 1:
        steps.tick(len(lists[0]))
        return list(lists[0]), steps.n

    result = lists[0]
    for other in lists[1:]:
        result = intersect(result, other, steps)
    return result, steps.n