"""Convert DBpedia/OpenAI parquet shards into float32 .npy arrays + metadata.

Corpus: the first 500,000 rows of the text-embedding-3-large-1536 1M set (shards 0-12).
Queries: 1,000 rows from shard 13, which is never loaded into any corpus.
3072-d: the 100K text-embedding-3-large-3072 set.
"""
import glob, json, os, sys
import numpy as np
import pyarrow.parquet as pq

DATA = sys.argv[1]
OUT = sys.argv[2]
os.makedirs(OUT, exist_ok=True)


def load(files, col, limit=None):
    vecs, ids, titles, texts = [], [], [], []
    n = 0
    for f in files:
        t = pq.read_table(f, columns=["_id", "title", "text", col])
        emb = t.column(col).combine_chunks()
        dim = len(emb[0].as_py())
        flat = emb.values.to_numpy(zero_copy_only=False).astype(np.float32)
        arr = flat.reshape(-1, dim)
        vecs.append(arr)
        ids += t.column("_id").to_pylist()
        titles += t.column("title").to_pylist()
        texts += t.column("text").to_pylist()
        n += len(arr)
        print(f, len(arr), dim, flush=True)
        if limit and n >= limit:
            break
    v = np.concatenate(vecs)
    if limit:
        v, ids, titles, texts = v[:limit], ids[:limit], titles[:limit], texts[:limit]
    return v, ids, titles, texts


c1536 = "text-embedding-3-large-1536-embedding"
files = sorted(glob.glob(f"{DATA}/d1536/*.parquet"))
v, ids, titles, texts = load(files, c1536, 500_000)
np.save(f"{OUT}/corpus1536.npy", v)
with open(f"{OUT}/corpus_meta.jsonl", "w") as fh:
    for i, t, x in zip(ids, titles, texts):
        fh.write(json.dumps({"id": i, "title": t, "text": x}) + "\n")
print("corpus", v.shape, "norms", np.linalg.norm(v[:1000], axis=1).min(), np.linalg.norm(v[:1000], axis=1).max())

q, qids, qtitles, qtexts = load(sorted(glob.glob(f"{DATA}/dq/*.parquet")), c1536)
rng = np.random.default_rng(42)
sel = np.sort(rng.choice(len(q), 1000, replace=False))
np.save(f"{OUT}/queries1536.npy", q[sel])
with open(f"{OUT}/queries_meta.jsonl", "w") as fh:
    for i in sel:
        fh.write(json.dumps({"id": qids[i], "title": qtitles[i]}) + "\n")
print("queries", q[sel].shape)

c3072 = "text-embedding-3-large-3072-embedding"
files = sorted(glob.glob(f"{DATA}/d3072/*.parquet"))
t0 = pq.read_table(files[0]).schema.names
col = [c for c in t0 if "embedding" in c][0]
v3, ids3, titles3, _ = load(files, col)
np.save(f"{OUT}/corpus3072.npy", v3)
with open(f"{OUT}/corpus3072_ids.json", "w") as fh:
    json.dump(ids3, fh)
print("3072", v3.shape, "overlap with 1536 corpus ids:", len(set(ids3) & set(ids)))