# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "sentence-transformers", "huggingface_hub"]
# ///
"""
Turn the harvest into journal centroid vectors for the recommender's stage 2.

    python build_centroids.py --db journals.db \
        --journal-text journal_text.jsonl --out centroids.npz

Run this AFTER harvest_openalex.py. It does no network I/O of its own beyond
downloading the sentence-transformer weights once.

What gets embedded
------------------
Per journal, in descending order of signal:

  1. OpenAlex topics, repeated in proportion to how much the journal publishes
     on each — an article-derived fingerprint, pre-aggregated by OpenAlex.
  2. The journal's ASJC subject names, from the index.
  3. The journal title.

Journals OpenAlex could not match still get (2) and (3), so every journal has a
vector and none silently drops out of the reranking.

Cost note
---------
all-MiniLM-L6-v2 is 384-dimensional and ~90 MB. 32,050 active journals × 384
float32 = ~49 MB uncompressed, well under what a Space can hold in memory. A
larger model buys little here: the input is a bag of topic labels, not prose,
so the extra capacity has nothing to work with.
"""
from __future__ import annotations

import argparse
import json
import sqlite3
from pathlib import Path

import numpy as np


def load_topics(path: str) -> dict[str, str]:
    out: dict[str, str] = {}
    p = Path(path)
    if not p.exists():
        print(f"note: {path} not found — building from index metadata only, "
              f"which is much weaker. Run harvest_openalex.py first.")
        return out
    with p.open(encoding="utf-8") as f:
        for line in f:
            try:
                rec = json.loads(line)
            except json.JSONDecodeError:
                continue
            sid = rec.get("scopus_id")
            if sid:
                out[sid] = (out.get(sid, "") + " " + rec.get("topics", "")).strip()
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", default="journals.db")
    ap.add_argument("--journal-text", default="journal_text.jsonl")
    ap.add_argument("--out", default="centroids.npz")
    ap.add_argument("--model", default="sentence-transformers/all-MiniLM-L6-v2")
    ap.add_argument("--batch-size", type=int, default=128)
    ap.add_argument("--active-only", action="store_true", default=True)
    ap.add_argument("--push-to-dataset", default=None)
    args = ap.parse_args()

    topics = load_topics(args.journal_text)
    con = sqlite3.connect(args.db)

    subs: dict[str, list[str]] = {}
    for sid, name in con.execute(
        """SELECT ja.scopus_id, coalesce(a.name, a.parent_name)
           FROM journal_asjc ja JOIN asjc a ON a.code = ja.asjc"""
    ):
        subs.setdefault(sid, []).append(name or "")

    where = "WHERE active = 1" if args.active_only else ""
    rows = con.execute(f"SELECT scopus_id, title FROM journals {where}").fetchall()
    con.close()

    ids, docs = [], []
    for sid, title in rows:
        parts = [title or ""]
        s = subs.get(sid)
        if s:
            parts.append(". ".join(sorted(set(s))))
        t = topics.get(sid)
        if t:
            parts.append(t)
        ids.append(sid)
        docs.append(". ".join(p for p in parts if p)[:2000])

    covered = sum(1 for i in ids if i in topics)
    print(f"{len(ids)} journals; {covered} ({covered/max(len(ids),1):.0%}) "
          f"have OpenAlex topics")

    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer(args.model)
    vecs = model.encode(
        docs, batch_size=args.batch_size, show_progress_bar=True,
        normalize_embeddings=True, convert_to_numpy=True,
    ).astype("float32")

    np.savez_compressed(args.out, ids=np.array(ids), vecs=vecs)
    size = Path(args.out).stat().st_size / 1e6
    print(f"wrote {args.out}  {vecs.shape}  ({size:.1f} MB)")

    if args.push_to_dataset:
        from huggingface_hub import HfApi
        api = HfApi()
        api.create_repo(args.push_to_dataset, repo_type="dataset", exist_ok=True)
        api.upload_file(path_or_fileobj=args.out, path_in_repo=Path(args.out).name,
                        repo_id=args.push_to_dataset, repo_type="dataset")
        print(f"pushed {args.out} -> {args.push_to_dataset}")


def load_centroids(path: str = "centroids.npz") -> dict[str, np.ndarray]:
    """Load into the {scopus_id: vector} mapping recommend() expects."""
    d = np.load(path, allow_pickle=False)
    return {str(i): v for i, v in zip(d["ids"], d["vecs"])}


if __name__ == "__main__":
    main()