#!/usr/bin/env python3
"""WP44 — normalise the filtered OpenAlex CS records (cs/*.jsonl.zst) into Parquet tables.

Tables (one Parquet file per input part under <out>/tables/<table>/<part>.parquet, resumable):
  works, authorships, work_topics, work_keywords, references, abstracts, locations, grants, counts_by_year
Run on the compute node after (or while) ingestion/openalex_snapshot.py runs; parts whose outputs exist are skipped.
"""
import argparse, glob, io, multiprocessing as mp, os, re, time, datetime as dt
import orjson, zstandard as zstd
import pyarrow as pa, pyarrow.parquet as pq

ARXIV_RE = re.compile(r"arxiv\.org/(?:abs|pdf)/([0-9]{4}\.[0-9]{4,5}(?:v\d+)?|[a-z\-]+(?:\.[A-Z]{2})?/[0-9]{7}(?:v\d+)?)", re.I)
ARXIV_DOI_RE = re.compile(r"^10\.48550/arxiv\.(.+)$", re.I)

S = {
 "works": pa.schema([("id", pa.int64()), ("doi", pa.string()), ("title", pa.string()), ("publication_year", pa.int32()),
    ("publication_date", pa.string()), ("created_date", pa.string()), ("updated_date", pa.string()), ("type", pa.string()),
    ("type_crossref", pa.string()), ("language", pa.string()), ("cited_by_count", pa.int32()), ("fwci", pa.float32()),
    ("citation_percentile", pa.float32()), ("is_top1", pa.bool_()), ("is_top10", pa.bool_()), ("is_oa", pa.bool_()), ("oa_status", pa.string()),
    ("license", pa.string()), ("is_retracted", pa.bool_()), ("is_paratext", pa.bool_()), ("is_xpac", pa.bool_()), ("n_refs", pa.int32()),
    ("n_authors", pa.int32()), ("n_institutions", pa.int32()), ("n_countries", pa.int32()), ("countries", pa.list_(pa.string())),
    ("institution_types", pa.list_(pa.string())), ("primary_topic_id", pa.int32()), ("primary_subfield_id", pa.int32()),
    ("primary_field_id", pa.int32()), ("primary_domain_id", pa.int32()), ("cs_primary", pa.bool_()), ("cs_any", pa.bool_()),
    ("source_id", pa.int64()), ("source_name", pa.string()), ("source_type", pa.string()), ("host_organization", pa.string()),
    ("is_in_doaj", pa.bool_()), ("arxiv_id", pa.string()), ("pmid", pa.string()), ("mag", pa.string()), ("has_abstract", pa.bool_()),
    ("has_fulltext", pa.bool_()), ("n_locations", pa.int32()), ("indexed_in", pa.list_(pa.string())), ("referenced_works_count", pa.int32()),
    ("n_grants", pa.int32()), ("apc_paid_usd", pa.int32()), ("is_authors_truncated", pa.bool_()), ("part", pa.string())]),
 "authorships": pa.schema([("work_id", pa.int64()), ("position", pa.int32()), ("author_position", pa.string()), ("author_id", pa.int64()),
    ("author_name", pa.string()), ("orcid", pa.string()), ("is_corresponding", pa.bool_()), ("institution_ids", pa.list_(pa.int64())),
    ("institution_names", pa.list_(pa.string())), ("institution_types", pa.list_(pa.string())), ("countries", pa.list_(pa.string())),
    ("n_raw_affiliations", pa.int32())]),
 "work_topics": pa.schema([("work_id", pa.int64()), ("rank", pa.int16()), ("topic_id", pa.int32()), ("score", pa.float32()),
    ("subfield_id", pa.int32()), ("field_id", pa.int32()), ("domain_id", pa.int32())]),
 "work_keywords": pa.schema([("work_id", pa.int64()), ("keyword", pa.string()), ("score", pa.float32())]),
 "references": pa.schema([("work_id", pa.int64()), ("ref_id", pa.int64())]),
 "abstracts": pa.schema([("work_id", pa.int64()), ("abstract", pa.string()), ("n_words", pa.int32())]),
 "locations": pa.schema([("work_id", pa.int64()), ("rank", pa.int16()), ("is_primary", pa.bool_()), ("is_best_oa", pa.bool_()), ("source_id", pa.int64()),
    ("source_name", pa.string()), ("source_type", pa.string()), ("landing_page_url", pa.string()), ("pdf_url", pa.string()), ("is_oa", pa.bool_()),
    ("license", pa.string()), ("version", pa.string()), ("is_accepted", pa.bool_()), ("is_published", pa.bool_())]),
 "grants": pa.schema([("work_id", pa.int64()), ("funder_id", pa.int64()), ("funder_name", pa.string()), ("award_id", pa.string())]),
 "counts_by_year": pa.schema([("work_id", pa.int64()), ("year", pa.int16()), ("cited_by_count", pa.int32())]),
}


def oid(u):
    if not u:
        return None
    t = u.rsplit("/", 1)[-1]
    i = 0
    while i < len(t) and not t[i].isdigit():
        i += 1
    try:
        return int(t[i:])
    except ValueError:
        return None


def abstract_text(inv):
    if not inv:
        return None
    pos = []
    for w, ps in inv.items():
        for p in ps:
            pos.append((p, w))
    pos.sort()
    return " ".join(w for _, w in pos)


def arxiv_from(w):
    ids = w.get("ids") or {}
    if ids.get("arxiv"):
        return str(ids["arxiv"]).rsplit("/", 1)[-1]
    doi = w.get("doi") or ""
    m = ARXIV_DOI_RE.match(doi.replace("https://doi.org/", ""))
    if m:
        return m.group(1)
    for loc in (w.get("locations") or []):
        for k in ("landing_page_url", "pdf_url"):
            u = loc.get(k) or ""
            m = ARXIV_RE.search(u)
            if m:
                return re.sub(r"v\d+$", "", m.group(1))
    return None


def process(args):
    path, out, cs_field = args
    part = os.path.basename(path).replace(".jsonl.zst", "")
    outs = {t: os.path.join(out, "tables", t, part + ".parquet") for t in S}
    if all(os.path.exists(p) for p in outs.values()):
        return part, "skip", 0
    for p in outs.values():
        os.makedirs(os.path.dirname(p), exist_ok=True)
    rows = {t: {c: [] for c in S[t].names} for t in S}
    n = 0
    with open(path, "rb") as fh, zstd.ZstdDecompressor().stream_reader(fh) as r:
        for line in io.TextIOWrapper(r, encoding="utf-8"):
            if not line.strip():
                continue
            w = orjson.loads(line); n += 1
            wid = oid(w["id"])
            doi = (w.get("doi") or "").replace("https://doi.org/", "").lower() or None
            pt = w.get("primary_topic") or {}
            topics = w.get("topics") or []
            cs_any = any(((t.get("field") or {}).get("id") == cs_field) for t in topics) or ((pt.get("field") or {}).get("id") == cs_field)
            loc = w.get("primary_location") or {}
            src = loc.get("source") or {}
            best = w.get("best_oa_location") or {}
            oa = w.get("open_access") or {}
            auths = w.get("authorships") or []
            inst_ids = set(); countries = set(); itypes = []
            for i, a in enumerate(auths):
                ins = a.get("institutions") or []
                ii = [oid(x.get("id")) for x in ins]
                inst_ids.update(x for x in ii if x)
                cc = [c for c in (a.get("countries") or []) if c]
                countries.update(cc)
                itypes.extend(x.get("type") for x in ins if x.get("type"))
                au = a.get("author") or {}
                R = rows["authorships"]
                R["work_id"].append(wid); R["position"].append(i); R["author_position"].append(a.get("author_position"))
                R["author_id"].append(oid(au.get("id"))); R["author_name"].append(au.get("display_name")); R["orcid"].append(au.get("orcid"))
                R["is_corresponding"].append(bool(a.get("is_corresponding"))); R["institution_ids"].append(ii)
                R["institution_names"].append([x.get("display_name") for x in ins]); R["institution_types"].append([x.get("type") for x in ins])
                R["countries"].append(cc); R["n_raw_affiliations"].append(len(a.get("raw_affiliation_strings") or []))
            for k, t in enumerate(topics):
                R = rows["work_topics"]
                R["work_id"].append(wid); R["rank"].append(k); R["topic_id"].append(oid(t.get("id"))); R["score"].append(t.get("score"))
                R["subfield_id"].append(oid((t.get("subfield") or {}).get("id"))); R["field_id"].append(oid((t.get("field") or {}).get("id")))
                R["domain_id"].append(oid((t.get("domain") or {}).get("id")))
            for kw in (w.get("keywords") or []):
                R = rows["work_keywords"]; R["work_id"].append(wid); R["keyword"].append(kw.get("display_name")); R["score"].append(kw.get("score"))
            refs = w.get("referenced_works") or []
            for rr in refs:
                rows["references"]["work_id"].append(wid); rows["references"]["ref_id"].append(oid(rr))
            ab = abstract_text(w.get("abstract_inverted_index"))
            if ab:
                R = rows["abstracts"]; R["work_id"].append(wid); R["abstract"].append(ab); R["n_words"].append(len(ab.split()))
            locs = w.get("locations") or []
            best_url = best.get("landing_page_url") if best else None
            for k, L in enumerate(locs):
                s_ = L.get("source") or {}
                R = rows["locations"]
                R["work_id"].append(wid); R["rank"].append(k); R["is_primary"].append(L.get("landing_page_url") == loc.get("landing_page_url") and L.get("pdf_url") == loc.get("pdf_url"))
                R["is_best_oa"].append(bool(best_url) and L.get("landing_page_url") == best_url)
                R["source_id"].append(oid(s_.get("id"))); R["source_name"].append(s_.get("display_name")); R["source_type"].append(s_.get("type"))
                R["landing_page_url"].append(L.get("landing_page_url")); R["pdf_url"].append(L.get("pdf_url")); R["is_oa"].append(bool(L.get("is_oa")))
                R["license"].append(L.get("license")); R["version"].append(L.get("version")); R["is_accepted"].append(bool(L.get("is_accepted"))); R["is_published"].append(bool(L.get("is_published")))
            for g in (w.get("grants") or []):
                R = rows["grants"]; R["work_id"].append(wid); R["funder_id"].append(oid(g.get("funder"))); R["funder_name"].append(g.get("funder_display_name")); R["award_id"].append(g.get("award_id"))
            for c in (w.get("counts_by_year") or []):
                R = rows["counts_by_year"]; R["work_id"].append(wid); R["year"].append(c.get("year")); R["cited_by_count"].append(c.get("cited_by_count"))
            cnp = w.get("citation_normalized_percentile") or {}
            ids = w.get("ids") or {}
            apc = (w.get("apc_paid") or {}).get("value_usd") if isinstance(w.get("apc_paid"), dict) else None
            R = rows["works"]
            vals = dict(id=wid, doi=doi, title=w.get("title") or w.get("display_name"), publication_year=w.get("publication_year"), publication_date=w.get("publication_date"),
                        created_date=w.get("created_date"), updated_date=(w.get("updated_date") or "")[:10] or None, type=w.get("type"), type_crossref=w.get("type_crossref"),
                        language=w.get("language"), cited_by_count=w.get("cited_by_count"), fwci=w.get("fwci"), citation_percentile=cnp.get("value"),
                        is_top1=bool(cnp.get("is_in_top_1_percent")), is_top10=bool(cnp.get("is_in_top_10_percent")), is_oa=bool(oa.get("is_oa")), oa_status=oa.get("oa_status"),
                        license=(best.get("license") if best else None) or loc.get("license"), is_retracted=bool(w.get("is_retracted")), is_paratext=bool(w.get("is_paratext")),
                        is_xpac=bool(w.get("is_xpac")), n_refs=len(refs), n_authors=w.get("authors_count") if w.get("authors_count") is not None else len(auths),
                        n_institutions=len(inst_ids), n_countries=len(countries), countries=sorted(countries), institution_types=sorted(set(itypes)),
                        primary_topic_id=oid(pt.get("id")), primary_subfield_id=oid((pt.get("subfield") or {}).get("id")), primary_field_id=oid((pt.get("field") or {}).get("id")),
                        primary_domain_id=oid((pt.get("domain") or {}).get("id")), cs_primary=(pt.get("field") or {}).get("id") == cs_field, cs_any=bool(cs_any),
                        source_id=oid(src.get("id")), source_name=src.get("display_name"), source_type=src.get("type"), host_organization=src.get("host_organization_name"),
                        is_in_doaj=bool(src.get("is_in_doaj")), arxiv_id=arxiv_from(w), pmid=(ids.get("pmid") or "").rsplit("/", 1)[-1] or None, mag=str(ids.get("mag")) if ids.get("mag") else None,
                        has_abstract=ab is not None, has_fulltext=bool(w.get("has_fulltext") or w.get("has_content")), n_locations=len(locs), indexed_in=w.get("indexed_in") or [],
                        referenced_works_count=w.get("referenced_works_count") if w.get("referenced_works_count") is not None else len(refs), n_grants=len(w.get("grants") or []),
                        apc_paid_usd=apc, is_authors_truncated=bool(w.get("is_authors_truncated")), part=part)
            for k, v in vals.items():
                R[k].append(v)
    for t, sch in S.items():
        tbl = pa.Table.from_pydict(rows[t], schema=sch)
        tmp = outs[t] + ".part"
        pq.write_table(tbl, tmp, compression="zstd"); os.replace(tmp, outs[t])
    return part, "done", n


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--cs-dir", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--workers", type=int, default=8); ap.add_argument("--limit", type=int, default=0)
    a = ap.parse_args()
    files = sorted(glob.glob(os.path.join(a.cs_dir, "*.jsonl.zst")))
    if a.limit:
        files = files[: a.limit]
    t0 = time.time(); tot = 0; done = 0
    with mp.Pool(a.workers, maxtasksperchild=40) as pool:
        for part, st, n in pool.imap_unordered(process, [(f, a.out, "https://openalex.org/fields/17") for f in files]):
            tot += n; done += 1
            if st == "done" or done % 100 == 0:
                print(f"[{dt.datetime.now():%H:%M:%S}] {st} {part} n={n} | {done}/{len(files)} files, {tot} works, {time.time()-t0:.0f}s", flush=True)
    print("finished", done, "files", tot, "works")


if __name__ == "__main__":
    main()