"""
Did retrieval find the right passage?  Steps 3 to 6.
====================================================
Data: SQuAD 1.1 development set (Rajpurkar et al., 2016; CC-BY-SA 4.0).
2,067 paragraphs from 48 Wikipedia articles, and 10,570 questions that people
wrote about them, each tied to the paragraph that answers it.
The search here is BM25, the standard keyword ranking used by most search
engines, written out in plain Python so you can read every line.

Setup (once):  nothing to install beyond Python 3
Run:           python retrieval_lesson.py                    (how often search finds the right paragraph)
               python retrieval_lesson.py overlap            (split by how many words the question borrows)
               python retrieval_lesson.py ask "your question" "expected answer"
                                                             (where the paragraph containing the answer ranks)
The first run downloads about 5 MB into ./data.
"""
import json
import math
import os
import re
import sys
import urllib.request
from collections import Counter

URL = 'https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json'
STOP = set("the a an of in on at to for and or is are was were be been by with as from that this which who whom "
           "what when where why how did does do it its his her their they he she".split())


def words(text):
    return [w for w in re.findall(r"\w+", text.lower()) if w not in STOP]


def load():
    os.makedirs('data', exist_ok=True)
    if not os.path.exists('data/squad_dev.json'):
        print('downloading SQuAD ...')
        urllib.request.urlretrieve(URL, 'data/squad_dev.json')
    d = json.load(open('data/squad_dev.json', encoding='utf-8'))
    paragraphs, questions = [], []
    for article in d['data']:
        for p in article['paragraphs']:
            pid = len(paragraphs)
            paragraphs.append((article['title'].replace('_', ' '), p['context']))
            for q in p['qas']:
                questions.append((q['question'], pid))
    return paragraphs, questions


class BM25:
    """Score = sum over query words of: rarity of the word x how often it appears in the paragraph,
    with diminishing returns for repeats (k1) and a penalty for long paragraphs (b)."""

    def __init__(self, paragraphs, k1=1.2, b=0.75):
        self.docs = [Counter(words(text)) for _, text in paragraphs]
        self.lengths = [sum(c.values()) for c in self.docs]
        self.avg = sum(self.lengths) / len(self.lengths)
        n = len(self.docs)
        df = Counter(w for c in self.docs for w in c)
        self.idf = {w: math.log(1 + (n - k + .5) / (k + .5)) for w, k in df.items()}
        self.index = {}
        for i, c in enumerate(self.docs):
            for w, f in c.items():
                self.index.setdefault(w, []).append((i, f))
        self.k1, self.b = k1, b

    def rank(self, query):
        scores = Counter()
        for w in set(words(query)):
            for i, f in self.index.get(w, []):
                norm = f + self.k1 * (1 - self.b + self.b * self.lengths[i] / self.avg)
                scores[i] += self.idf[w] * f * (self.k1 + 1) / norm
        return [i for i, _ in scores.most_common()]      # paragraphs sharing no word are never returned


def position(ranking, pid):
    return ranking.index(pid) + 1 if pid in ranking else None     # None = not returned at all


def report(paragraphs, questions, bm):
    pos = [position(bm.rank(q), pid) for q, pid in questions]
    rate = lambda k: sum(p is not None and p <= k for p in pos) / len(pos)
    mrr = sum(1 / p for p in pos if p) / len(pos)
    print(f"BM25 keyword search over {len(paragraphs):,} paragraphs, {len(questions):,} questions\n")
    for k in (1, 3, 5, 10, 20):
        print(f"  right paragraph in the top {k:<2}  {rate(k):.0%}")
    print(f"  mean reciprocal rank       {mrr:.2f}")
    print(f"  not returned at all        {sum(p is None for p in pos) / len(pos):.1%}")
    wrong = [(q, pid) for (q, pid), p in zip(questions, pos) if p != 1]
    same = sum(paragraphs[bm.rank(q)[0]][0] == paragraphs[pid][0] for q, pid in wrong if bm.rank(q))
    print(f"\nWhen the top result is wrong, it comes from the same article {same / len(wrong):.0%} of the time.")


def overlap(paragraphs, questions, bm):
    groups = {'a third or less': [], 'about half': [], 'two thirds or more': []}
    for q, pid in questions:
        qw = set(words(q))
        share = len(qw & bm.docs[pid].keys()) / max(1, len(qw))
        key = 'a third or less' if share < .34 else 'about half' if share < .67 else 'two thirds or more'
        groups[key].append(position(bm.rank(q), pid))
    print("Share of the question's words that also appear in its paragraph, and how search does\n")
    print(f"  {'words shared':<20} {'questions':>10} {'top 1':>7} {'top 5':>7}")
    for key, pos in groups.items():
        r = lambda k: sum(p is not None and p <= k for p in pos) / len(pos)
        print(f"  {key:<20} {len(pos) / len(questions):>10.0%} {r(1):>7.0%} {r(5):>7.0%}")


def ask(paragraphs, bm, query, answer):
    ranking = bm.rank(query)
    hits = [i for i in ranking if answer.lower() in paragraphs[i][1].lower()]
    print(f'Query: "{query}"\n')
    for n, i in enumerate(ranking[:3], 1):
        print(f"  {n}. [{paragraphs[i][0]}] {paragraphs[i][1][:150]}...")
    first = ranking.index(hits[0]) + 1 if hits else None
    print(f'\nFirst paragraph containing "{answer}": ' + (f"rank {first}" if first else "not returned at all"))


if __name__ == '__main__':
    paragraphs, questions = load()
    bm = BM25(paragraphs)
    if len(sys.argv) >= 4 and sys.argv[1] == 'ask':
        ask(paragraphs, bm, sys.argv[2], sys.argv[3])
    elif len(sys.argv) >= 2 and sys.argv[1] == 'overlap':
        overlap(paragraphs, questions, bm)
    else:
        report(paragraphs, questions, bm)