"""Memory-lean data prep using polars directly on the hub parquet shards."""
import glob
import os
import sys

import numpy as np
import polars as pl

N_TRAIN_Q = 80_000
N_DEV_Q   = 2_000
N_NEGS    = 4
SEED      = 42

SNAP = glob.glob(os.path.expanduser(
    "~/.cache/huggingface/hub/datasets--sentence-transformers--msmarco/snapshots/*"))[0]

print("scanning margin-mse shards...", flush=True)
df = (
    pl.scan_parquet(f"{SNAP}/bert-ensemble-margin-mse/*.parquet")
    .select(
        pl.col("query_id").cast(pl.Int64),
        pl.col("positive_id").cast(pl.Int64),
        pl.col("negative_id").cast(pl.Int64),
        pl.col("score").cast(pl.Float32),
    )
    .collect()
)
print("rows:", df.height, flush=True)

# canonical positive per query = first occurrence; then 4 distinct negatives for it
first_pos = df.group_by("query_id").agg(pl.col("positive_id").first().alias("pos"))
g = (
    df.join(first_pos, on="query_id")
    .filter((pl.col("positive_id") == pl.col("pos")) & (pl.col("negative_id") != pl.col("pos")))
    .unique(subset=["query_id", "negative_id"], keep="first", maintain_order=True)
    .group_by("query_id")
    .agg(
        pl.col("pos").first(),
        pl.col("negative_id").head(N_NEGS).alias("negs"),
        pl.col("score").head(N_NEGS).alias("margins"),
    )
    .filter(pl.col("negs").list.len() >= N_NEGS)
)
print("eligible queries:", g.height, flush=True)

g = g.sample(n=min(N_TRAIN_Q + N_DEV_Q, g.height), seed=SEED, shuffle=True)
print("sampled:", g.height, flush=True)

qids = g["query_id"].to_numpy()
pos = g["pos"].to_numpy()
negs = np.stack(g["negs"].to_numpy())            # [N, 4]
margins = np.stack(g["margins"].to_numpy()).astype(np.float32)

need_docs = np.unique(np.concatenate([pos, negs.reshape(-1)]))
print("unique docs:", len(need_docs), "queries:", len(qids), flush=True)

print("loading corpus/queries parquet...", flush=True)
corpus_files = glob.glob(f"{SNAP}/corpus/*.parquet")
queries_files = glob.glob(f"{SNAP}/queries/*.parquet")
assert corpus_files and queries_files, "run: hf download for corpus/queries first"

corpus = (
    pl.scan_parquet(corpus_files)
    .select(pl.col("passage_id").cast(pl.Int64).alias("id"), pl.col("passage").alias("text"))
    .filter(pl.col("id").is_in(need_docs))
    .collect()
)
queries = (
    pl.scan_parquet(queries_files)
    .select(pl.col("query_id").cast(pl.Int64).alias("id"), pl.col("query").alias("text"))
    .filter(pl.col("id").is_in(qids))
    .collect()
)
print("resolved docs:", corpus.height, "queries:", queries.height, flush=True)
assert corpus.height == len(need_docs), "missing corpus texts"
assert queries.height == len(qids), "missing query texts"

qtext = dict(zip(queries["id"].to_list(), queries["text"].to_list()))
dtext = dict(zip(corpus["id"].to_list(), corpus["text"].to_list()))

clean = lambda s: " ".join(s.split())
id2idx = {}
with open("/home/anon/pog/data/manifest.tsv", "w") as f:
    for q in qids:
        id2idx[f"q{q}"] = len(id2idx)
        f.write(f"q{q}\tq\t{clean(qtext[q])}\n")
    for d in need_docs:
        id2idx[f"d{d}"] = len(id2idx)
        f.write(f"d{d}\td\t{clean(dtext[d])}\n")
print("manifest rows:", len(id2idx), flush=True)

q_idx = np.array([id2idx[f"q{q}"] for q in qids], dtype=np.int64)
p_idx = np.array([id2idx[f"d{d}"] for d in pos], dtype=np.int64)
n_idx = np.array([[id2idx[f"d{d}"] for d in row] for row in negs], dtype=np.int64)

tr = slice(0, N_TRAIN_Q)
dv = slice(N_TRAIN_Q, None)
np.savez("/home/anon/pog/data/train.npz", q=q_idx[tr], pos=p_idx[tr], neg=n_idx[tr], margin=margins[tr])
np.savez("/home/anon/pog/data/dev.npz",   q=q_idx[dv], pos=p_idx[dv], neg=n_idx[dv], margin=margins[dv])
print(f"DONE train={N_TRAIN_Q} dev={len(qids)-N_TRAIN_Q}", flush=True)