"""d01_acc_pass.py (device side). One pass over every acc.csv (imi_vm20i only).

For each file: parse with the pandas C engine, check the header and the time step,
compute per-channel statistics, write a float32 cache (device scratch, not delivered)
and a 10 Hz block-RMS envelope, and write one JSON of statistics per run.
Resumable: runs already done are skipped; stops before the time budget (argv[1], s).
"""
import sys, os, json, time, glob
import numpy as np, pandas as pd

ROOT = os.path.expanduser("~/mnt/archive (1)")
OUT = os.path.join(ROOT, "Phase2", "derived", "acc_stats")
CACHE = os.path.expanduser("~/work/acc_cache")
FS_DOC = 51200.0  # documented sampling rate (Hz), used only for comparison
budget = float(sys.argv[1]) if len(sys.argv) > 1 else 140.0
t_start = time.time()
os.makedirs(OUT, exist_ok=True); os.makedirs(CACHE, exist_ok=True)
files = sorted(glob.glob(os.path.join(ROOT, "imi_vm20i", "dataset*", "*_acc.csv")),
               key=lambda p: int(p.split(os.sep)[-2].replace("dataset", "")))
done = 0
for f in files:
    run = f.split(os.sep)[-2]
    tag = f"imi_vm20i_{run}"
    js = os.path.join(OUT, tag + ".json")
    if os.path.exists(js):
        continue
    size = os.path.getsize(f)
    if time.time() - t_start + size / 80e6 + 5 > budget:
        break
    t0 = time.time()
    with open(f, "r", encoding="utf-8", errors="strict") as fh:
        header = fh.readline().strip()
    df = pd.read_csv(f, dtype={"time": np.float64, "acc0": np.float64, "acc1": np.float64}, engine="c")
    t = df["time"].to_numpy(); a = df[["acc0", "acc1"]].to_numpy()
    n = len(t)
    d = np.diff(t)
    med = float(np.median(d))
    rec = {
        "machine": "imi_vm20i", "run": run, "file": os.path.basename(f), "bytes": size,
        "header": header, "columns": list(df.columns), "rows": int(n),
        "nan_time": int(np.isnan(t).sum()), "nan_acc0": int(np.isnan(a[:, 0]).sum()), "nan_acc1": int(np.isnan(a[:, 1]).sum()),
        "t_first": float(t[0]), "t_last": float(t[-1]),
        "dt_median": med, "dt_min": float(d.min()), "dt_max": float(d.max()),
        "n_dt_off_1e-9": int((np.abs(d - med) > 1e-9).sum()),
        "n_dt_nonpositive": int((d <= 0).sum()),
        "fs_est": 1.0 / med, "fs_doc": FS_DOC,
        "duration_s": float(t[-1] - t[0] + med),
    }
    for k in (0, 1):
        x = a[:, k]
        mx, mn = float(np.nanmax(x)), float(np.nanmin(x))
        rec[f"acc{k}"] = {
            "mean": float(np.nanmean(x)), "std": float(np.nanstd(x)), "min": mn, "max": mx,
            "n_at_max": int((x == mx).sum()), "n_at_min": int((x == mn).sum()),
            "p0001": float(np.nanpercentile(x, 0.01)), "p9999": float(np.nanpercentile(x, 99.99)),
        }
    # float32 cache (device scratch only)
    np.save(os.path.join(CACHE, tag + ".npy"), a.astype(np.float32))
    # 10 Hz block RMS envelope (5120 samples per block at 51.2 kHz), mean removed per block
    B = int(round(1.0 / med / 10.0))
    m = (n // B) * B
    blk = a[:m].reshape(-1, B, 2)
    env = np.sqrt(((blk - blk.mean(axis=1, keepdims=True)) ** 2).mean(axis=1)).astype(np.float32)
    np.save(os.path.join(OUT, tag + "_env10hz.npy"), env)
    rec["env_block_samples"] = B; rec["env_rows"] = int(env.shape[0])
    rec["seconds_to_process"] = round(time.time() - t0, 1)
    with open(js, "w") as fh:
        json.dump(rec, fh, indent=1)
    done += 1
    del df, t, a, d, blk, env
remaining = sum(1 for f in files if not os.path.exists(os.path.join(OUT, f"imi_vm20i_{f.split(os.sep)[-2]}.json")))
print(f"processed {done} files this call; remaining {remaining} of {len(files)}; elapsed {time.time()-t_start:.0f}s")