#!/usr/bin/env python3
"""Offset detector v3 — pre-trigger floor, with pre/mid/end consistency.

Brian's method, and better than v2's whole-record median for one reason: the
pre-trigger window is *definitionally* quiet (it is the buffer captured before
the trigger fired), whereas a whole-record median is merely robust to the event.

Per channel:
    pre  = median of the first `pretrig_samples` samples (STRT record)
    mid  = median of the middle third
    end  = median of the final third
    spread = max(pre,mid,end) - min(pre,mid,end)

A DC offset is a *constant floor*: |pre| at or above the floor AND a small
spread. A transient (settling, handling, a long-tailed event) moves one segment
relative to the others and is rejected by the spread test.

Floor default 0.025 in/s = 5 A/D counts (Instantel's own criterion; 1 count =
0.005 in/s).  Quantisation is 0.005 in/s, so `spread` is measured in units of it.
"""
from __future__ import annotations
import argparse, csv, re, statistics, sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from minimateplus.event_file_io import read_blastware_file

GEO=("Tran","Vert","Long"); K=10.0/32000.0
_WAVE=re.compile(r"\.[A-Za-z0-9]{2}0[Ww]$"); _STEM=re.compile(r"^([B-Z])(\d{3})")

def serial_of(n):
    m=_STEM.match(n)
    return f"BE{(ord(m.group(1))-ord('B'))*1000+int(m.group(2))}" if m else "?"

def scan(ps):
    p=Path(ps)
    try:
        ev=read_blastware_file(p); s=ev.raw_samples or {}
        if not all(s.get(c) for c in GEO): return None
        pre_n=ev.pretrig_samples
        ts=ev.timestamp
        stamp=(f"{ts.year:04d}-{ts.month:02d}-{ts.day:02d}T"
               f"{ts.hour:02d}:{ts.minute:02d}:{ts.second:02d}") if ts else ""
        out=[]
        for ch in GEO:
            a=s[ch]; n=len(a); t=n//3
            pre = a[:pre_n] if (pre_n and 0 < pre_n < n) else a[:t]
            mid, end = a[t:2*t], a[2*t:]
            if not pre or not mid or not end: continue
            v=[statistics.median(x)*K for x in (pre,mid,end)]
            out.append({"serial":serial_of(p.name),"timestamp":stamp,
                        "filename":p.name,"channel":ch,
                        "pretrig_n": pre_n or 0,
                        "pre":round(v[0],4),"mid":round(v[1],4),"end":round(v[2],4),
                        "spread":round(max(v)-min(v),4),
                        "peak":round(max(abs(x) for x in a)*K,4)})
        return out
    except Exception:
        return None

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--dir",required=True); ap.add_argument("--jobs",type=int,default=4)
    ap.add_argument("--floor",type=float,default=0.025)
    ap.add_argument("--max-spread",type=float,default=0.02)
    ap.add_argument("--out",required=True)
    a=ap.parse_args()
    seen=set(); files=[]
    for q in sorted(Path(a.dir).rglob("*")):
        if q.is_file() and _WAVE.search(q.name) and q.name not in seen:
            seen.add(q.name); files.append(str(q))
    print(f"unique waveform binaries: {len(files)}",flush=True)
    rows=[]
    with ProcessPoolExecutor(max_workers=a.jobs) as ex:
        for i,f in enumerate(as_completed([ex.submit(scan,p) for p in files]),1):
            r=f.result()
            if r: rows.extend(r)
            if i%2000==0: print(f"  {i}/{len(files)}",flush=True)
    for r in rows:
        r["offset"]=int(abs(r["pre"])>=a.floor and r["spread"]<=a.max_spread)
    cols=["serial","timestamp","filename","channel","pretrig_n","pre","mid","end","spread","peak","offset"]
    with open(a.out,"w",newline="") as fh:
        w=csv.DictWriter(fh,fieldnames=cols); w.writeheader(); w.writerows(rows)
    from collections import defaultdict
    per=defaultdict(set); tot=defaultdict(set)
    for r in rows:
        tot[r["serial"]].add(r["filename"])
        if r["offset"]: per[r["serial"]].add(r["filename"])
    print(f"\nfloor={a.floor} in/s ({a.floor/0.005:.0f} counts)  max spread={a.max_spread}")
    print(f"units affected: {len(per)} of {len(tot)}")
    for s in sorted(per,key=lambda s:-len(per[s])):
        print(f"  {s:9} {len(per[s]):4} / {len(tot[s]):4} events")
    print(f"\nwrote {a.out}")

if __name__=="__main__": main()