"""d03_segment_features.py (device side). Segment-level features for every labelled cutting path.

For each run and each high-rate channel (Acc0 and Acc1 on imi_vm20i from the d01 float32 cache;
Sound0, Sound1, Sound2 from the WAV files) and each label.csv row:
  time-domain features of the mean-removed path signal: RMS, peak, crest factor, excess kurtosis,
  share of samples at 16-bit full scale (sound only), number of samples;
  Welch power spectral density (Hann window, nfft 8192, 50 percent overlap), bins 0 to 2048
  (0 to 12.8 kHz for 51.2 kHz accelerometers, 0 to 12 kHz for 48 kHz sound);
  the same for a pre-cut window [start - 1.2 s, start - 0.2 s] when it lies after the previous
  path end plus 0.2 s and after time zero (PSD with nfft 2048, bins 0 to 512).
Times follow the article: CSV time columns are seconds on the common timeline, WAV sample k is at
k / 48000 s. Resumable, one output file per run; time budget in argv[1] (s).
"""
import sys, os, json, time, glob, wave
import numpy as np

ROOT = os.path.expanduser("~/mnt/archive (1)")
OUT = os.path.join(ROOT, "Phase2", "derived", "seg_features")
CACHE = os.path.expanduser("~/work/acc_cache")
ACCSTAT = os.path.join(ROOT, "Phase2", "derived", "acc_stats")
budget = float(sys.argv[1]) if len(sys.argv) > 1 else 150.0
t_start = time.time()
NFFT, NB = 8192, 2049
NFFT_PRE, NB_PRE = 2048, 513
PRE0, PRE1, GAP = 1.2, 0.2, 0.2


def welch(x, fs, nfft, nb):
    w = np.hanning(nfft).astype(np.float32)
    U = float((w.astype(np.float64) ** 2).sum())
    hop = nfft // 2
    if len(x) < nfft:
        x = np.pad(x, (0, nfft - len(x)))
    view = np.lib.stride_tricks.sliding_window_view(x, nfft)[::hop]
    acc = np.zeros(nfft // 2 + 1, dtype=np.float64)
    n = 0
    for i in range(0, view.shape[0], 256):
        fr = view[i:i + 256] * w
        X = np.fft.rfft(fr, axis=1)
        acc += (X.real ** 2 + X.imag ** 2).sum(axis=0)
        n += fr.shape[0]
    P = acc / n / (fs * U)
    P[1:-1] *= 2.0
    return P[:nb].astype(np.float32), n


def tfeat(x, clip_raw=None):
    x = x.astype(np.float64)
    x = x - x.mean()
    m2 = float((x ** 2).mean())
    rms = np.sqrt(m2)
    pk = float(np.abs(x).max())
    kurt = float((x ** 4).mean() / m2 ** 2 - 3.0) if m2 > 0 else float("nan")
    clip = float(np.mean((clip_raw == 32767) | (clip_raw == -32768))) if clip_raw is not None else float("nan")
    return [rms, pk, pk / rms if rms > 0 else float("nan"), kurt, clip, float(len(x))]


def runkey(p):
    return (p.split(os.sep)[-2], int(p.split(os.sep)[-1].replace("dataset", "")))


runs = sorted([d for m in ("imi_vm20i", "imi_vmx30ui", "tmf_vf10") for d in glob.glob(os.path.join(ROOT, m, "dataset*"))], key=runkey)
done = 0
for d in runs:
    mach, run = d.split(os.sep)[-2], d.split(os.sep)[-1]
    out = os.path.join(OUT, f"{mach}_{run}.npz")
    if os.path.exists(out):
        continue
    est = 20.0 if mach != "tmf_vf10" else 45.0
    if time.time() - t_start + est > budget:
        break
    t0 = time.time()
    lab = np.genfromtxt(os.path.join(d, "label.csv"), delimiter=",", names=True, dtype=None, encoding="utf-8")
    starts = np.atleast_1d(lab["start"]).astype(float); ends = np.atleast_1d(lab["end"]).astype(float)
    order = np.argsort(starts)
    prev_end = np.full(len(starts), -np.inf)
    se = sorted(zip(starts, ends))
    for i in range(len(starts)):
        # latest end among paths that start before this one
        prev = [e for s, e in se if s < starts[i]]
        prev_end[i] = max(prev) if prev else -np.inf
    chans = {}
    tag = f"{mach}_{run}"
    if os.path.exists(os.path.join(CACHE, tag + ".npy")):
        a = np.load(os.path.join(CACHE, tag + ".npy"), mmap_mode="r")
        st = json.load(open(os.path.join(ACCSTAT, tag + ".json")))
        fs_a = 1.0 / st["dt_median"]; t0a = st["t_first"]
        chans["acc0"] = ("acc", a[:, 0], fs_a, t0a, None)
        chans["acc1"] = ("acc", a[:, 1], fs_a, t0a, None)
    for f in sorted(glob.glob(os.path.join(d, "*.wav"))):
        ch = os.path.basename(f).rsplit("_", 1)[-1].replace(".wav", "")
        with wave.open(f, "rb") as w:
            fs_w = w.getframerate(); raw = np.frombuffer(w.readframes(w.getnframes()), dtype="<i2")
        chans[ch] = ("snd", raw, float(fs_w), 0.0, raw)
    res = {"start": starts, "end": ends, "label": np.atleast_1d(lab["label"]).astype(int)}
    meta = {}
    for ch, (typ, sig, fs, tz, raw) in chans.items():
        F = np.full((len(starts), 6), np.nan); P = np.full((len(starts), NB), np.nan, dtype=np.float32)
        FP = np.full((len(starts), 6), np.nan); PP = np.full((len(starts), NB_PRE), np.nan, dtype=np.float32)
        nfr = np.zeros(len(starts), dtype=np.int32)
        for i, (s, e) in enumerate(zip(starts, ends)):
            i0 = int(round((s - tz) * fs)); i1 = int(round((e - tz) * fs))
            i0 = max(i0, 0); i1 = min(i1, len(sig))
            if i1 - i0 < 64:
                continue
            x = np.asarray(sig[i0:i1], dtype=np.float32)
            F[i] = tfeat(x, None if raw is None else raw[i0:i1])
            P[i], nfr[i] = welch(x - x.mean(), fs, NFFT, NB)
            p0, p1 = s - PRE0, s - PRE1
            if p0 >= 0 and p0 >= prev_end[i] + GAP:
                j0 = int(round((p0 - tz) * fs)); j1 = int(round((p1 - tz) * fs))
                if j1 <= len(sig) and j1 - j0 > NFFT_PRE:
                    y = np.asarray(sig[j0:j1], dtype=np.float32)
                    FP[i] = tfeat(y, None if raw is None else raw[j0:j1])
                    PP[i], _ = welch(y - y.mean(), fs, NFFT_PRE, NB_PRE)
        res[f"{ch}_feat"] = F; res[f"{ch}_psd"] = P; res[f"{ch}_pre_feat"] = FP; res[f"{ch}_pre_psd"] = PP; res[f"{ch}_nframes"] = nfr
        meta[ch] = {"fs": fs, "df": fs / NFFT, "df_pre": fs / NFFT_PRE, "type": typ}
    res["meta"] = np.array(json.dumps({"machine": mach, "run": run, "channels": meta,
                                        "feat_cols": ["rms", "peak", "crest", "kurt_excess", "fullscale_share", "n_samples"],
                                        "nfft": NFFT, "nb": NB, "nfft_pre": NFFT_PRE, "nb_pre": NB_PRE,
                                        "pre_window": [PRE0, PRE1, GAP], "seconds": round(time.time() - t0, 1)}))
    np.savez(out, **res)
    done += 1
left = sum(1 for d in runs if not os.path.exists(os.path.join(OUT, f"{d.split(os.sep)[-2]}_{d.split(os.sep)[-1]}.npz")))
print(f"segment features: {done} runs this call; remaining {left} of {len(runs)}; elapsed {time.time()-t_start:.0f}s")