"""Train + save the memorization classifiers used by the demo's your-own-text mode.

Mirrors score_1m_nodup.py: train on the Run-1 deduplicated features with base-model-
reproduced positives removed, 1:3 rebalanced, on the six no-dup features. One classifier
per dataset (they transfer, per the cross-domain result, but we keep all four). Saves the
fitted StandardScaler + GradientBoostingClassifier per dataset via joblib.
"""
import os, numpy as np, pandas as pd, joblib
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler

ROOT = "/projects/nulab/borkar.j/predicting_memorization"
OUT = f"{ROOT}/demo_app/models"
FEATS = ["zlib_entropy", "base_ppl", "base_loss_variance",
         "gradient_norm", "one_step_abs_change", "one_step_rel_change"]
DATASETS = {"fineweb": "FineWeb", "pg19": "PG-19", "code": "The Stack", "owm": "OpenWebMath"}
os.makedirs(OUT, exist_ok=True)

for d, disp in DATASETS.items():
    df = pd.read_csv(f"{ROOT}/{d}_features.csv")
    m = np.load(f"{ROOT}/base_extractability/base_memorized_{d}.npy")
    pi = df.index[df.label == 1].tolist()
    df["bm"] = 0
    df.loc[pi, "bm"] = m.astype(int)
    df = df[df.approx_dup_count == 1]                       # dedup
    df = df[~((df.label == 1) & (df.bm == 1))]             # remove base-reproduced positives
    pos = df[df.label == 1]; neg = df[df.label == 0]; n = 3 * len(pos)
    if len(neg) > n:
        neg = neg.sample(n=n, random_state=42)
    tr = pd.concat([pos, neg]).sample(frac=1, random_state=42).reset_index(drop=True)
    X = np.nan_to_num(tr[FEATS].values, posinf=1e10, neginf=-1e10)
    y = tr.label.values
    sc = StandardScaler().fit(X)
    clf = GradientBoostingClassifier(n_estimators=200, max_depth=5, learning_rate=0.1,
                                     subsample=0.8, random_state=42).fit(sc.transform(X), y)
    joblib.dump({"scaler": sc, "clf": clf, "feats": FEATS, "dataset": disp},
                f"{OUT}/clf_{d}.joblib")
    print(f"{disp:12s} trained on {len(tr)} rows ({int(y.sum())} pos); saved clf_{d}.joblib")
print("done")