from jev import ask, fits


def question(key):
    return {
        "type": "noul",
        "instructions": f"Does the passage `passages.{key}` help answer `query`?",
        "criteria": {"true": "The passage states facts, numbers, or claims that directly answer the query.",
                     "false": "The passage is off topic or only shares words with the query."},
    }


def batches(query, candidates):
    batch = {}
    for i, chunk in enumerate(candidates):
        key = f"p{i}"
        trial = dict(batch, **{key: chunk})
        state = {"query": query, "passages": {k: passage(c) for k, c in trial.items()}}
        if batch and not fits(state, {k: question(k) for k in trial}):
            yield batch
            trial = {key: chunk}
        batch = trial
    if batch:
        yield batch


def passage(chunk):
    return {"title": chunk["title"], "section": chunk["section"], "text": chunk["text"]}


def rerank(query, candidates, keep=5):
    scored = []
    for batch in batches(query, candidates):
        state = {"query": query, "passages": {k: passage(c) for k, c in batch.items()}}
        answers = ask(state, {k: question(k) for k in batch}, stage="rerank")
        scored += [dict(c, relevance=answers[k]["noul"]) for k, c in batch.items()]
    scored.sort(key=lambda c: (-c["relevance"], -c["similarity"]))
    return scored[:keep]