#!/usr/bin/env python3
"""WP44 — embed title + abstract of corpus works with a scientific sentence-embedding model, sharded for multi-node runs.

  --prepare   : DuckDB exports data/processed/embeddings/input/shard_NNN.parquet (work_id, text) for eligible works with an
                abstract (or title only when --title-only-fallback), N shards of --shard-rows rows.
  --run       : encode shards [--shard-from, --shard-to) with --model on MPS/CPU → data/processed/embeddings/<model_tag>/
                shard_NNN.npy (float16, L2-normalised) + shard_NNN.ids.npy. Resumable (existing shards skipped).
Provenance: data/processed/embeddings/<model_tag>/manifest.json (model, revision, max_seq_length, host, timings).
"""
import argparse, glob, json, os, socket, time, datetime as dt
import numpy as np


def prepare(root, shard_rows, title_only_fallback):
    import duckdb
    con = duckdb.connect(os.path.join(root, "data/processed/wp44.duckdb"), read_only=True); con.execute("SET threads TO 16;")
    out = os.path.join(root, "data/processed/embeddings/input"); os.makedirs(out, exist_ok=True)
    fb = " OR TRUE" if title_only_fallback else ""
    con.execute(f"""CREATE OR REPLACE TEMP TABLE emb_in AS
        SELECT w.id AS work_id, w.title || '. ' || coalesce(substr(a.abstract, 1, 2000), '') AS text
        FROM corpus_works w LEFT JOIN oa_abstracts a ON a.work_id = w.id
        WHERE w.eligible AND (a.abstract IS NOT NULL{fb}) ORDER BY w.id""")
    n = con.execute("SELECT count(*) FROM emb_in").fetchone()[0]
    con.execute(f"COPY (SELECT *, (row_number() OVER (ORDER BY work_id) - 1) // {shard_rows} AS shard FROM emb_in) TO '{out}' (FORMAT PARQUET, PARTITION_BY (shard), OVERWRITE_OR_IGNORE)")
    print("prepared", n, "texts in", len(glob.glob(out + "/shard=*")), "shards")


def run(root, model_name, shard_from, shard_to, batch, max_len):
    import torch, pyarrow.parquet as pq
    from sentence_transformers import SentenceTransformer
    tag = model_name.replace("/", "__"); out = os.path.join(root, "data/processed/embeddings", tag); os.makedirs(out, exist_ok=True)
    dev = "mps" if torch.backends.mps.is_available() else "cpu"
    model = SentenceTransformer(model_name, device=dev); model.max_seq_length = max_len
    shards = sorted(glob.glob(os.path.join(root, "data/processed/embeddings/input", "shard=*", "*.parquet")))
    man_p = os.path.join(out, f"manifest_{socket.gethostname()}.json"); man = json.load(open(man_p)) if os.path.exists(man_p) else {"model": model_name, "max_seq_length": max_len, "device": dev, "host": socket.gethostname(), "shards": {}}
    for p in shards:
        k = int(os.path.basename(os.path.dirname(p)).split("=")[1])
        if not (shard_from <= k < shard_to):
            continue
        o = os.path.join(out, f"shard_{k:04d}.npy")
        if os.path.exists(o):
            continue
        t = pq.read_table(p).to_pydict(); t0 = time.time()
        E = model.encode(t["text"], batch_size=batch, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False).astype(np.float16)
        np.save(o + ".part.npy", E); os.replace(o + ".part.npy", o); np.save(o.replace(".npy", ".ids.npy"), np.asarray(t["work_id"], dtype=np.int64))
        man["shards"][str(k)] = {"n": len(t["text"]), "seconds": round(time.time() - t0, 1), "when": dt.datetime.now(dt.UTC).isoformat()}
        json.dump(man, open(man_p, "w"), indent=1)
        print(f"shard {k}: {len(t['text'])} texts in {time.time()-t0:.0f}s ({len(t['text'])/(time.time()-t0):.0f}/s)", flush=True)


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--prepare", action="store_true"); ap.add_argument("--run", action="store_true")
    ap.add_argument("--model", default="allenai/specter2_base"); ap.add_argument("--shard-rows", type=int, default=50000); ap.add_argument("--shard-from", type=int, default=0)
    ap.add_argument("--shard-to", type=int, default=10**6); ap.add_argument("--batch", type=int, default=128); ap.add_argument("--max-len", type=int, default=256); ap.add_argument("--title-only-fallback", action="store_true")
    a = ap.parse_args()
    if a.prepare:
        prepare(a.root, a.shard_rows, a.title_only_fallback)
    if a.run:
        run(a.root, a.model, a.shard_from, a.shard_to, a.batch, a.max_len)


if __name__ == "__main__":
    main()