"""
Can an automatic score tell a good summary from a bad one?  Steps 3 to 6.
=========================================================================
Data: SummEval (Fabbri et al., 2021), as packaged by MTEB (mteb/summeval, MIT).
100 CNN/DailyMail news articles. For each: 16 summaries written by different
summarization systems, and 11 summaries written by people. Experts rated every
machine summary from 1 to 5 on four things:
  relevance    does it capture the important points?
  coherence    is it well organized?
  fluency      is each sentence well written?
  consistency  does it only state things the article supports?

Setup (once):  pip install pandas pyarrow
Run:           python rouge_lesson.py            (how well ROUGE agrees with the experts)
               python rouge_lesson.py cases      (summaries ROUGE loves and the experts don't)
               python rouge_lesson.py bestof     (use ROUGE to pick the best of 16 summaries)
The first run downloads the data (about 0.4 MB) into ./data.
"""
import os
import re
import sys
import urllib.request
from collections import Counter

import pandas as pd

URL = ('https://huggingface.co/datasets/mteb/summeval/resolve/main/data/'
       'test-00000-of-00001-35901af5f6649399.parquet')
DIMS = ['relevance', 'coherence', 'fluency', 'consistency']


def words(text):
    return re.findall(r"\w+", text.lower())      # \w keeps accented letters like the u in Yuki


def rouge1(summary, reference):
    """ROUGE-1 F1: overlap of single words, each word counted at most as often as it appears in both."""
    s, r = Counter(words(summary)), Counter(words(reference))
    overlap = sum((s & r).values())
    if overlap == 0:
        return 0.0
    precision = overlap / sum(s.values())         # share of the summary's words found in the reference
    recall = overlap / sum(r.values())            # share of the reference's words found in the summary
    return 2 * precision * recall / (precision + recall)


def rouge2(summary, reference):
    """Same idea with pairs of adjacent words."""
    pairs = lambda t: Counter(zip(t, t[1:]))
    s, r = pairs(words(summary)), pairs(words(reference))
    overlap = sum((s & r).values())
    if overlap == 0:
        return 0.0
    p, rc = overlap / sum(s.values()), overlap / sum(r.values())
    return 2 * p * rc / (p + rc)


def load():
    os.makedirs('data', exist_ok=True)
    if not os.path.exists('data/summeval.parquet'):
        print('downloading SummEval ...')
        urllib.request.urlretrieve(URL, 'data/summeval.parquet')
    d = pd.read_parquet('data/summeval.parquet')
    rows = []
    for article, a in d.iterrows():
        for system, summary in enumerate(a.machine_summaries):
            row = {'article': article, 'system': system, 'summary': summary, 'words': len(words(summary)),
                   # score against every human summary and keep the best match, the usual practice
                   'rouge1': max(rouge1(summary, h) for h in a.human_summaries),
                   'rouge2': max(rouge2(summary, h) for h in a.human_summaries)}
            for dim in DIMS:
                row[dim] = a[dim][system]
            rows.append(row)
    x = pd.DataFrame(rows)
    x['expert'] = x[DIMS].mean(axis=1)
    return d, x


def agreement(x):
    cols = ['rouge1', 'rouge2', 'words']
    print("How well each score agrees with the experts (Spearman correlation: 1 = same order, 0 = no relation)")
    print("\nOne summary at a time (all 1,600 summaries)")
    print(x[cols + DIMS].corr(method='spearman').loc[cols, DIMS].round(2).to_string())
    sysl = x.groupby('system')[cols + DIMS].mean()
    print("\nWhole systems (each system's average over 100 articles, 16 systems)")
    print(sysl.corr(method='spearman').loc[cols, DIMS].round(2).to_string())
    print(f"\nShare of summaries the experts rated 3 or lower on consistency: {(x.consistency <= 3).mean():.0%}")


def cases(d, x, limit=3):
    x = x.copy()
    x['rank'] = x.groupby('article').rouge1.rank(ascending=False, method='first')
    bad = x[(x['rank'] <= 3) & (x.consistency <= 2)].sort_values('rouge1', ascending=False)
    print(f"{len(bad)} summaries are in their article's top 3 by ROUGE-1 but rated 2 or lower on consistency.\n")
    for _, r in bad.head(limit).iterrows():
        print("=" * 78)
        print(f"article {r.article}: ROUGE-1 {r.rouge1:.2f} (#{int(r['rank'])} of 16), consistency {r.consistency:.2f}, relevance {r.relevance:.2f}")
        print("\nSUMMARY:", r.summary)
        print("\nARTICLE (start):", d.loc[r.article, 'text'][:600], "...")


def bestof(x):
    pick = lambda col: x.loc[x.groupby('article')[col].idxmax()]
    print("Pick one summary per article out of 16, four different ways. Average expert ratings of the picks:\n")
    print(f"  {'how the pick was made':<28} {'relevance':>9} {'consistency':>11} {'picks rated 3 or lower on consistency':>38}")
    for label, p in [('at random (average of all)', x), ('highest ROUGE-1', pick('rouge1')),
                     ('longest summary', pick('words')), ('experts\' favorite', pick('expert'))]:
        print(f"  {label:<28} {p.relevance.mean():>9.2f} {p.consistency.mean():>11.2f} {(p.consistency <= 3).mean():>38.0%}")


if __name__ == '__main__':
    d, x = load()
    if len(sys.argv) > 1 and sys.argv[1] == 'cases':
        cases(d, x)
    elif len(sys.argv) > 1 and sys.argv[1] == 'bestof':
        bestof(x)
    else:
        agreement(x)