import json, os
import pyarrow.parquet as pq

os.chdir("/home/corov/cyber/data")
meta = json.load(open("cgym_meta.json"))
names = [s["rfilename"] for s in meta["siblings"]]
tops = sorted(set(n.split("/")[1] for n in names if n.startswith("data/")))
print("data subdirs:", tops)

tr = pq.read_table("elfsong/train/train-00000-of-00001.parquet").to_pandas()
ev = pq.read_table("elfsong/eval/test-00000-of-00001.parquet").to_pandas()
BL = {"42536536", "42537493", "42537664", "42537686", "42537734", "42538131",
      "383170474", "383825645"}
tr["tid"] = tr.apply(lambda r: f'{r["kind"]}:{r["id"]}', axis=1)
ev["tid"] = ev.apply(lambda r: f'{r["kind"]}:{r["id"]}', axis=1)
print("train rows:", len(tr), "| eval rows:", len(ev))
tr_excluded = tr[tr["id"].astype(str).isin(BL)]
print("excluded in train:", len(tr_excluded), tr_excluded["tid"].tolist())
tr_ok = tr[~tr["id"].astype(str).isin(BL)]
print("usable train:", len(tr_ok))
tr_ok.to_json("train_ids.json", orient="records")
ev.to_json("eval_ids.json", orient="records")
overlap = set(tr_ok["tid"]) & set(ev["tid"])
print("TRAIN-EVAL OVERLAP:", len(overlap))

for kind in ["arvo", "oss-fuzz"]:
    pats = [n for n in names if f"/{kind}/" in n]
    print(kind, "files in cgym:", len(pats), pats[:3])

# which of our needed files exist
need_train = {f"data/{k}/{i}/description.txt" for k, i in zip(tr_ok["kind"], tr_ok["id"].astype(str))}
have = set(names)
missing = [f for f in need_train if f not in have]
print("train description.txt missing:", len(missing), missing[:5])

# project diversity
print("train project counts (top 15):")
print(tr_ok["project"].value_counts().head(15).to_string())
print("train kind counts:", tr_ok["kind"].value_counts().to_dict())