# 500k rows × 12 features ≈ 48 MB parquet. Enough to trip the 10 MB
# max_committed_size_mb warning below — dial it up to 5_000_000 for a
# "really big" run, down to 50_000 for a fast smoke test. The downstream
# cache-key tracks this value, so the subgraph re-runs on any change.
N_ROWS = 500_000
N_FEATURES = 12
SEED = 42

# Generate the "big" parquet. Not hand-written data — a reproducible
# synthetic sample derived from the seed so any reviewer gets identical
# bytes without checking in the file.
#
# Reads `data/seed.json` directly via `jc.load` inside this cell (rather
# than through a separate `seed` cell) so the subgraph works cleanly
# across mixed cache-hit / re-run boundaries — the seed config doesn't
# need to live in the kernel's memory.
import numpy as np
import pandas as pd

import jellycell.api as jc

seed_cfg = jc.load("data/seed.json")
rng = np.random.default_rng(SEED + seed_cfg["offset"])
X = rng.normal(loc=seed_cfg["mean"], scale=seed_cfg["scale"], size=(N_ROWS, N_FEATURES))
df = pd.DataFrame(X, columns=[f"f{i}" for i in range(N_FEATURES)])
df["label"] = (df["f0"] + 0.5 * df["f1"] - 0.25 * df["f2"] > 0).astype(int)

# Explicit path keeps the .gitignore glob simple.
out = jc.save(df, "artifacts/large_data/sample_dataset.parquet")
print(f"{len(df):,} rows × {df.shape[1]} cols → {out}")

# A compact digest of the underlying data. This IS committed (few KB) and
# carries enough summary stats for the tearsheet. It's a good pattern:
# commit the summary, git-ignore the bulk.
#
# Reads both the parquet bytes and its own file size off disk so this cell
# works when `sample` is cached (the in-memory `out` Path from `sample`
# isn't available across cache/run boundaries).
from pathlib import Path as _Path

parquet = _Path("artifacts/large_data/sample_dataset.parquet")
df = jc.load(str(parquet))
headline = {
    "rows": len(df),
    "features": df.shape[1] - 1,
    "positive_rate": round(float(df["label"].mean()), 4),
    "feature_mean": round(float(df.drop(columns=["label"]).values.mean()), 4),
    "feature_std": round(float(df.drop(columns=["label"]).values.std()), 4),
    "size_mb": round(parquet.stat().st_size / (1024 * 1024), 2),
}
jc.save(headline, "artifacts/large_data/headline.json")
print(headline)

# A tiny figure that fits in a tearsheet. Path-less — layout=by_notebook
# from jellycell.toml drops it under artifacts/large_data/ automatically.
import matplotlib.pyplot as plt

counts = df["label"].value_counts().sort_index()
fig, ax = plt.subplots(figsize=(5, 3))
ax.bar(["0", "1"], counts.values, color=["#6b7280", "#4f46e5"])
ax.set_ylabel("Count")
ax.set_title(f"Label balance (n={len(df):,})")
ax.grid(alpha=0.3, axis="y")
for i, v in enumerate(counts.values):
    ax.text(i, v, f"{v:,}", ha="center", va="bottom", fontsize=9)
fig.tight_layout()
jc.figure(fig=fig)