"""Assemble the precomputed lookup table for the demo.

For each (Pythia, dataset) Run-2 set, join token sequences + memorized labels +
classifier score + base-model features, decode the 50-token prefix/suffix to text,
and keep a compact sample (all memorized up to a cap, plus a random sample of
non-memorized). Public / non-PII datasets only.

Output: demo_app/data/lookup.parquet
"""
import os, numpy as np, pandas as pd
from transformers import AutoTokenizer

ROOT = "/projects/nulab/borkar.j/predicting_memorization"
TOK = f"{ROOT}/checkpoints/pythia-1.4b"
DATASETS = {"fineweb": "FineWeb", "pg19": "PG-19", "code": "The Stack", "owm": "OpenWebMath"}
PREFIX_LEN, WINDOW_LEN = 50, 100
MEM_CAP, NONMEM_N, SEED = 3000, 5000, 42
OUT = f"{ROOT}/demo_app/data"

tok = AutoTokenizer.from_pretrained(TOK)
rng = np.random.default_rng(SEED)
rows = []

for d, disp in DATASETS.items():
    ck = f"{ROOT}/checkpoints/pythia_{d}_run2_seed42"
    ds = np.load(f"{ck}/dataset.npy")                        # [1M,128]
    lab = np.load(f"{ck}/memorization_labels.npy").astype(bool)
    fe = pd.read_csv(f"{ROOT}/full_run2_results/features/{d}_run2_full_features.csv")
    sc = pd.read_csv(f"{ROOT}/full_run2_results/cleaned/scores/{d}_run2_scores.csv")
    assert len(ds) == len(lab) == len(fe) == len(sc), d

    mem = np.where(lab)[0]
    non = np.where(~lab)[0]
    if len(mem) > MEM_CAP:
        mem = rng.choice(mem, MEM_CAP, replace=False)
    non = rng.choice(non, min(NONMEM_N, len(non)), replace=False)
    idx = np.concatenate([mem, non])

    pref = tok.batch_decode(ds[idx, :PREFIX_LEN].tolist(), skip_special_tokens=True)
    suf = tok.batch_decode(ds[idx, PREFIX_LEN:WINDOW_LEN].tolist(), skip_special_tokens=True)
    sub = pd.DataFrame({
        "model": "Pythia-1.4B", "dataset": disp,
        "prefix": pref, "suffix": suf,
        "memorized": lab[idx],
        "base_reproduced": sc["base_reproduced"].values[idx].astype(bool),
        "score": sc["score"].values[idx],
        "zlib_entropy": fe["zlib_entropy"].values[idx],
        "base_ppl": fe["base_ppl"].values[idx],
        "base_loss_variance": fe["base_loss_variance"].values[idx],
        "gradient_norm": fe["gradient_norm"].values[idx],
        "one_step_rel_change": fe["one_step_rel_change"].values[idx],
    })
    rows.append(sub)
    print(f"{disp:12s} kept {len(sub)} rows ({int(lab[idx].sum())} memorized)")

os.makedirs(OUT, exist_ok=True)
df = pd.concat(rows, ignore_index=True)
df.to_parquet(f"{OUT}/lookup.parquet", index=False)
print(f"\nwrote {OUT}/lookup.parquet  ({len(df)} rows, {df.memory_usage(deep=True).sum()/1e6:.1f} MB in mem)")