#!/usr/bin/env python3
"""WP44 — semantic k-NN edges between works (Phase 3, semantic layer) with FAISS HNSW on the SPECTER2 embeddings.
Edges: (work_id, neighbour_id, cosine_similarity, rank) for k = --k neighbours, similarity >= --min-sim; provenance = model tag +
index parameters; verification_status = 'machine-verified' (deterministic from embeddings). Also computes, per work, the share
of its k neighbours that belong to a different primary category (semantic boundary-crossing, used in RQ5) and writes
results/graph/semantic_crossing_yearly.csv. Output: data/processed/graph/semantic_edges_<tag>.parquet."""
import argparse, glob, json, os, time, numpy as np, pandas as pd


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--model", default="allenai/specter2_base"); ap.add_argument("--k", type=int, default=10)
    ap.add_argument("--min-sim", type=float, default=0.80); ap.add_argument("--sample", type=int, default=0, help="query only a random sample of works (0 = all)")
    a = ap.parse_args(); R = a.root; tag = a.model.replace("/", "__"); out = f"{R}/data/processed/graph"; os.makedirs(out, exist_ok=True); os.makedirs(f"{R}/results/graph", exist_ok=True)
    import faiss, duckdb
    E, I = [], []
    for f in sorted(glob.glob(f"{R}/data/processed/embeddings/{tag}/shard_*.npy")):
        if f.endswith(".ids.npy"): continue
        E.append(np.load(f)); I.append(np.load(f.replace(".npy", ".ids.npy")))
    X = np.concatenate(E).astype(np.float32); ids = np.concatenate(I); faiss.normalize_L2(X); n, d = X.shape; print("vectors", n, d, flush=True)
    t0 = time.time(); index = faiss.IndexHNSWFlat(d, 32, faiss.METRIC_INNER_PRODUCT); index.hnsw.efConstruction = 80; index.add(X); print(f"index built in {time.time()-t0:.0f}s", flush=True)
    index.hnsw.efSearch = 64
    q_idx = np.arange(n) if not a.sample else np.random.default_rng(44).choice(n, size=min(a.sample, n), replace=False)
    rows = []; B = 50000
    for s in range(0, len(q_idx), B):
        qi = q_idx[s:s + B]; D, N = index.search(X[qi], a.k + 1)
        for r, (qq, dd, nn) in enumerate(zip(qi, D, N)):
            rk = 0
            for sim, j in zip(dd, nn):
                if j == qq or j < 0: continue
                rk += 1
                if sim >= a.min_sim:
                    rows.append((int(ids[qq]), int(ids[j]), float(sim), rk))
        print(f"{s+len(qi)}/{len(q_idx)} queried, {len(rows)} edges, {time.time()-t0:.0f}s", flush=True)
    df = pd.DataFrame(rows, columns=["work_id", "neighbour_id", "cosine", "rank"]); df.to_parquet(f"{out}/semantic_edges_{tag}.parquet")
    json.dump({"model": a.model, "k": a.k, "min_sim": a.min_sim, "n_vectors": int(n), "n_edges": int(len(df)), "index": "HNSW32 efC80 efS64 inner-product", "seconds": round(time.time() - t0)}, open(f"{out}/semantic_edges_{tag}.json", "w"), indent=1)
    con = duckdb.connect(f"{R}/data/processed/wp44.duckdb", read_only=True)
    con.execute(f"CREATE OR REPLACE VIEW work_primary AS SELECT * FROM read_parquet('{R}/data/processed/classification/primary/*.parquet')")
    con.execute(f"CREATE OR REPLACE VIEW sem AS SELECT * FROM read_parquet('{out}/semantic_edges_{tag}.parquet')")
    con.execute(f"""COPY (SELECT c.publication_year AS year, pa.primary_category, count(*) n_edges, round(avg((pa.primary_category <> pb.primary_category)::int),4) share_cross_category
                 FROM sem s JOIN work_primary pa ON pa.work_id = s.work_id JOIN work_primary pb ON pb.work_id = s.neighbour_id JOIN corpus_works c ON c.id = s.work_id
                 GROUP BY 1,2 ORDER BY 1,2) TO '{R}/results/graph/semantic_crossing_yearly.csv' (HEADER)""")
    print("done", len(df), "edges")


if __name__ == "__main__":
    main()