"""Phase B — deduplication of discovery records and stage-1 (title/abstract) rule-based screening.
Inputs: data/metadata/s2_records.jsonl (+ openalex_records.jsonl if present). Outputs: data/metadata/records_dedup.csv,
data/metadata/stage1_candidates.csv, data/metadata/screening_stage1_rules.csv, reports/search_audit.md (counts section).
"""
import csv, json, re, sys, unicodedata
from collections import Counter
from pathlib import Path
import pandas as pd, yaml

ROOT = Path(__file__).resolve().parents[2]; META = ROOT / "data" / "metadata"


def norm_title(t):
    t = unicodedata.normalize("NFKD", t or "").encode("ascii", "ignore").decode().lower()
    return re.sub(r"[^a-z0-9]", "", t)


def load():
    rows = []
    f = META / "s2_records.jsonl"
    if f.exists():
        for line in open(f, encoding="utf-8"):
            p = json.loads(line)
            ext = p.get("externalIds") or {}
            rows.append({"source": p.get("_source"), "query": p.get("_query"), "sid": p.get("paperId"), "doi": (ext.get("DOI") or "").lower(),
                         "arxiv": ext.get("ArXiv"), "title": p.get("title") or "", "year": p.get("year"), "venue": (p.get("venue") or (p.get("publicationVenue") or {}).get("name") or ""),
                         "authors": "; ".join(a.get("name", "") for a in (p.get("authors") or [])), "abstract": p.get("abstract") or "",
                         "citations": p.get("citationCount"), "oa_pdf": ((p.get("openAccessPdf") or {}).get("url") or ""), "is_oa": p.get("isOpenAccess"),
                         "pub_types": ";".join(p.get("publicationTypes") or []), "fields": ";".join(p.get("fieldsOfStudy") or []), "pub_date": p.get("publicationDate")})
    f2 = META / "openalex_records.jsonl"
    if f2.exists():
        for line in open(f2, encoding="utf-8"):
            w = json.loads(line)
            rows.append({"source": "openalex", "query": w.get("_query"), "sid": w.get("id"), "doi": (w.get("doi") or "").replace("https://doi.org/", "").lower(),
                         "arxiv": None, "title": w.get("title") or "", "year": w.get("publication_year"), "venue": ((w.get("primary_location") or {}).get("source") or {}).get("display_name", ""),
                         "authors": "; ".join((a.get("author") or {}).get("display_name", "") for a in w.get("authorships") or []), "abstract": w.get("abstract") or "",
                         "citations": w.get("cited_by_count"), "oa_pdf": ((w.get("best_oa_location") or {}).get("pdf_url") or ""), "is_oa": (w.get("open_access") or {}).get("is_oa"),
                         "pub_types": w.get("type"), "fields": "", "pub_date": w.get("publication_date")})
    f3 = META / "crossref_records.jsonl"
    if f3.exists():
        for line in open(f3, encoding="utf-8"):
            it = json.loads(line)
            yr = ((it.get("issued") or {}).get("date-parts") or [[None]])[0][0]
            abst = re.sub(r"<[^>]+>", " ", it.get("abstract") or ""); abst = re.sub(r"\s+", " ", abst).strip()
            rows.append({"source": "crossref", "query": it.get("_query"), "sid": None, "doi": (it.get("DOI") or "").lower(), "arxiv": None,
                         "title": " ".join(it.get("title") or []), "year": yr, "venue": " ".join(it.get("container-title") or []),
                         "authors": "; ".join(f"{a.get('given','')} {a.get('family','')}".strip() for a in it.get("author") or []), "abstract": abst,
                         "citations": it.get("is-referenced-by-count"), "oa_pdf": "", "is_oa": None, "pub_types": it.get("type"), "fields": ";".join(it.get("subject") or []), "pub_date": None})
    f4 = META / "arxiv_records.jsonl"
    if f4.exists():
        for line in open(f4, encoding="utf-8"):
            it = json.loads(line); aid = it["id"].split("/abs/")[-1]
            rows.append({"source": "arxiv", "query": it.get("_query"), "sid": None, "doi": (it.get("doi") or "").lower(), "arxiv": re.sub(r"v\d+$", "", aid),
                         "title": it["title"], "year": int(it["published"][:4]), "venue": "arXiv", "authors": "; ".join(it["authors"]), "abstract": it["abstract"],
                         "citations": None, "oa_pdf": f"https://arxiv.org/pdf/{aid}", "is_oa": True, "pub_types": "preprint", "fields": ";".join(it["categories"]), "pub_date": it["published"][:10]})
    df = pd.DataFrame(rows)
    df["record_id"] = [r.sid if isinstance(r.sid, str) and r.sid else ("doi_" + re.sub(r"[^a-z0-9]", "_", r.doi) if r.doi else ("arxiv_" + str(r.arxiv) if r.arxiv else f"row_{i}")) for i, r in df.iterrows()]
    return df


def dedup(df):
    df = df.copy(); df["ntitle"] = df.title.map(norm_title); df["first_author"] = df.authors.str.split(";").str[0].str.strip().str.split().str[-1].str.lower()
    df["key_doi"] = df.doi.where(df.doi != "", None)
    df["key_title"] = df.ntitle + "|" + df.year.fillna(0).astype(int).astype(str)
    df["raw_n"] = 1
    # prefer records with abstract, DOI, OA pdf
    df["score_rec"] = df.abstract.str.len().fillna(0) + 1000 * (df.doi != "") + 500 * (df.oa_pdf != "")
    df = df.sort_values("score_rec", ascending=False)
    # group by DOI first, then by normalised title (+year within 1 handled by exact title match)
    seen_doi = {}; seen_title = {}; keep = []; queries = {}
    for i, r in df.iterrows():
        k = None
        if r.key_doi and r.key_doi in seen_doi:
            k = seen_doi[r.key_doi]
        elif r.ntitle and len(r.ntitle) > 15 and r.ntitle in seen_title:
            k = seen_title[r.ntitle]
        if k is None:
            k = i; keep.append(i)
            if r.key_doi: seen_doi[r.key_doi] = k
            if r.ntitle: seen_title[r.ntitle] = k
        queries.setdefault(k, set()).add(str(r.query)); 
    out = df.loc[keep].copy(); out["n_queries"] = out.index.map(lambda k: len(queries[k])); out["queries"] = out.index.map(lambda k: " || ".join(sorted(queries[k])))
    return out.reset_index(drop=True)


def screen(df, kw):
    text = (df.title.fillna("") + " . " + df.abstract.fillna("")).str.lower()
    def hits(group):
        pats = [k.lower() for k in kw[group]]
        return text.map(lambda t: sum(1 for p_ in pats if p_ in t))
    df = df.copy()
    df["h_hedonic"] = hits("hedonic_strong"); df["h_price"] = hits("price_outcome"); df["h_res"] = hits("residential"); df["h_reg"] = hits("regression_signal"); df["h_excl"] = hits("exclude_hard")
    df["score"] = 3 * (df.h_hedonic > 0) + 2 * (df.h_price > 0) + 1 * (df.h_res > 0) + 1 * (df.h_reg > 0) - 5 * ((df.h_excl > 0) & (df.h_price == 0))
    df["stage1_candidate"] = ((df.h_price > 0) & (df.h_res > 0) & ((df.h_hedonic > 0) | (df.h_reg > 0))) | ((df.h_hedonic > 0) & (df.h_res > 0))
    df["stage1_reason"] = df.apply(lambda r: ("candidate" if r.stage1_candidate else ("no residential price outcome" if r.h_price == 0 else "no hedonic/regression signal")), axis=1)
    df["no_abstract"] = df.abstract.fillna("").str.len() < 50
    return df


def main():
    kw = yaml.safe_load((ROOT / "scripts" / "screening" / "keywords.yaml").read_text())
    raw = load(); dd = dedup(raw); sc = screen(dd, kw)
    sc.to_csv(META / "records_dedup.csv", index=False)
    cand = sc[sc.stage1_candidate].sort_values(["score", "citations"], ascending=False)
    cand.to_csv(META / "stage1_candidates.csv", index=False)
    sc[["sid", "doi", "title", "year", "score", "stage1_candidate", "stage1_reason", "no_abstract"]].to_csv(META / "screening_stage1_rules.csv", index=False)
    summary = {"raw_records": int(len(raw)), "unique_records": int(len(dd)), "with_doi": int((dd.doi != "").sum()), "with_abstract": int((~sc.no_abstract).sum()),
               "stage1_candidates": int(len(cand)), "candidates_with_oa_pdf": int((cand.oa_pdf != "").sum()), "by_source": raw.source.value_counts().to_dict(),
               "years": {"min": int(dd.year.min()), "max": int(dd.year.max())}}
    (META / "screening_summary.json").write_text(json.dumps(summary, indent=2)); print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()