#!/usr/bin/env python3
"""Probe: readout-vocabulary coverage curve (Zipf knee) to pick vocab size.

Replicates the exact EPT fit/calibrate/segment pipeline of SPIRA.build on
train.bin, counts patch frequencies, then measures on held-out streams
(test.bin, oneshot.bin) the fraction of target patches covered by the
top-V vocabulary for V in a grid. Writes JSON to data/vocab_probe.json.
"""
import json
import sys
import time
from collections import Counter

sys.path.insert(0, "/home/z/my-project/download/spira")

import numpy as np

from spira.ept import EPT
from spira.hashing import fnv1a64

BASE = "/home/z/my-project"

t0 = time.perf_counter()
train = open(f"{BASE}/data/train.bin", "rb").read()
test = open(f"{BASE}/data/test.bin", "rb").read()
oneshot = open(f"{BASE}/data/oneshot.bin", "rb").read()

ept = EPT(budget_bits=16.0, min_len=2, max_len=32, alpha=0.5)
ept.fit(train[:2_000_000])
budget = ept.calibrate(train[2_000_000:2_400_000], target_len=4.5)
print(f"EPT budget={budget:.2f} bits ({time.perf_counter()-t0:.0f}s)")

def ids_of(data: bytes) -> np.ndarray:
    patches = ept.segment(data)
    return np.fromiter((fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches))

ids_train = ids_of(train)
ids_test = ids_of(test)
ids_one = ids_of(oneshot)
print(f"patches: train={len(ids_train)} test={len(ids_test)} oneshot={len(ids_one)} "
      f"({time.perf_counter()-t0:.0f}s)")

uniq, cnt = np.unique(ids_train, return_counts=True)
# same deterministic tie-break as SPIRA.build: count desc, larger id first
order = np.lexsort((uniq, -cnt))
sorted_ids = uniq[order]

GRID = [4096, 6144, 8192, 12288, 16384, 24576, 32768]
rows = []
for V in GRID:
    vocab = set(sorted_ids[:V].tolist())
    cov_test = float(np.mean([int(i) in vocab for i in ids_test.tolist()]))
    cov_one = float(np.mean([int(i) in vocab for i in ids_one.tolist()]))
    mass_train = float(cnt[order][:V].sum() / cnt.sum())
    rows.append({"vocab": V, "cov_test": round(cov_test, 4),
                 "cov_oneshot": round(cov_one, 4),
                 "train_mass": round(mass_train, 4)})
    print(f"V={V:>6}  held-out test coverage={cov_test:.4f}  "
          f"oneshot coverage={cov_one:.4f}  train mass={mass_train:.4f}")

out = {"budget_bits": float(budget), "distinct_train": int(len(uniq)),
       "n_train": int(len(ids_train)), "curve": rows,
       "elapsed_s": round(time.perf_counter() - t0, 1)}
with open(f"{BASE}/data/vocab_probe.json", "w") as f:
    json.dump(out, f, indent=1)
print("saved -> data/vocab_probe.json")