"""Classify papers as surveys or research papers with SciGenre.

    python predict.py papers.csv out.csv [--model REPO_ID_OR_PATH]

papers.csv needs the columns `title` and `abstract`. The output adds `p_survey` and `category` (survey, research,
or no abstract when the abstract has fewer than ten words, which the model does not classify).

The model input is model_text(title, abstract): the title and the abstract with markup, LaTeX and publisher
boilerplate removed, exactly as in training. The model never sees the venue.
"""
import argparse
import html
import re

import numpy as np
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

THRESHOLD, MAX_LEN, MIN_WORDS, SURVEY = 0.5, 384, 10, 0


def clean_text(text) -> str:
    """Remove markup, LaTeX and formatting artifacts; collapse whitespace."""
    if text is None or (isinstance(text, float) and text != text):  # None / NaN
        return ""
    text = str(text)
    text = re.sub(r"<tex-math.*?</tex-math>", " ", text, flags=re.DOTALL)  # JATS math
    text = re.sub(r"<[^>]+>", " ", text)                                # HTML / JATS tags
    text = re.sub(r"\$[^$]*\$", " ", text)                             # inline LaTeX math
    text = re.sub(r"\\[a-zA-Z]+\*?(?:\{[^}]*\})?", " ", text)          # LaTeX commands
    text = re.sub(r"[{}\[\]]", " ", text)                              # brackets
    text = re.sub(r"^\s*abstract[:.\s-]*", "", text, flags=re.IGNORECASE)
    return re.sub(r"\s+", " ", text).strip()


# Publisher and indexing boilerplate that some abstracts carry (copyright lines, licences, journal citation
# lines, WIREs category lists, Annual Reviews notices). It names publishers and venues, so the model never sees it.
BOILERPLATE = re.compile(
    r"(?:©|\(c\)\s*(?:19|20)\d\d|\bcopyright\s*(?:©|\(c\)|(?:19|20)\d\d)|\ball rights reserved\b|"
    r"\bthis article is categorized under\b|\bexpected final online publication date\b|\bother information published in\b|"
    r"\blicen[cs]e:\s*https?://|\bpublished online in wiley\b|\bcreative commons attribution\b)",
    re.IGNORECASE,
)


# Headings of structured abstracts ("Background:", "Methods:", ...) at the start of a sentence.
SECTION_HEADER = re.compile(
    r"(?:^|(?<=[.!?] ))(?:background|objectives?|purpose|aims?|methods?|methodology|design(?:/methodology/approach)?|"
    r"approach|results?|findings|conclusions?|originality(?:/value)?|value|significance|context|introduction|"
    r"motivation|limitations|implications)\s*:\s*",
    re.IGNORECASE,
)


def strip_boilerplate(text: str) -> str:
    """Decode HTML entities, drop structured-abstract headings, then cut trailing boilerplate (marker in the last
    60% of the text) or drop the sentence with the marker."""
    text = SECTION_HEADER.sub("", html.unescape(text))
    text = re.sub(r"\bdoi:\s*10\.\S+", " ", text, flags=re.IGNORECASE)
    m = BOILERPLATE.search(text)
    while m:
        if m.start() >= 0.4 * len(text):
            text = text[:m.start()]
        else:
            s = text.rfind(". ", 0, m.start()) + 1
            e = text.find(". ", m.end()); e = len(text) if e < 0 else e + 1
            text = text[:s] + text[e:]
        m = BOILERPLATE.search(text)
    return re.sub(r"\s+", " ", text).strip()


def model_text(title, abstract) -> str:
    """Input of the final model: title and abstract only, boilerplate removed (never the venue)."""
    return f"{clean_text(title)}. {strip_boilerplate(clean_text(abstract))}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("papers"); ap.add_argument("out"); ap.add_argument("--model", default=".")
    ap.add_argument("--batch", type=int, default=32)
    a = ap.parse_args()
    df = pd.read_csv(a.papers)
    tok = AutoTokenizer.from_pretrained(a.model)
    model = AutoModelForSequenceClassification.from_pretrained(a.model, dtype=torch.float32).eval()
    ok = df["abstract"].fillna("").map(lambda s: len(clean_text(s).split()) >= MIN_WORDS).values
    texts = [model_text(t, s) for t, s in zip(df["title"].fillna(""), df["abstract"].fillna(""))]
    p = np.full(len(df), np.nan)
    idx = np.where(ok)[0]
    with torch.no_grad():
        for i in range(0, len(idx), a.batch):
            b = idx[i:i + a.batch]
            enc = tok([texts[j] for j in b], truncation=True, max_length=MAX_LEN, padding=True, return_tensors="pt")
            p[b] = torch.softmax(model(**enc).logits.float(), -1)[:, SURVEY].numpy()
    df["p_survey"] = p
    df["category"] = np.where(~ok, "no abstract", np.where(p >= THRESHOLD, "survey", "research"))
    df.to_csv(a.out, index=False)
    print(df["category"].value_counts().to_string())


if __name__ == "__main__":
    main()