"""p03_stats.py: exploratory statistics for the pre-specified descriptive questions D2 to D5, D7, D10.

Every interval is a 95 percent percentile interval from a bootstrap that resamples runs (dataset
folders) with replacement within each machine, so paths from one run (same tool, workpiece and
setup) move together. Statistics are written as sums of per-run pieces so that each resample is a
weighted sum (resample counts times run pieces), which gives exactly the statistic recomputed on the
resampled runs. All estimates are exploratory; nothing is confirmatory in Phase 2.
"""
import json, sys, glob
import numpy as np, pandas as pd
sys.path.insert(0, "/home/claude/mill/code")
from vals import Values

OUT = "/home/claude/mill/outputs"
RAW = "/home/claude/mill/raw"
B = 4000
SEED = 20260922
LEVEL = 0.95
V = Values(script="p03_stats.py")
T = pd.read_parquet(f"{OUT}/segments.parquet")
LAB = ("imi_vm20i", "imi_vmx30ui")
res = {"B": B, "seed": SEED}
V.set("boot_b", B, "resamples", {"population": "bootstrap setting", "filter": "all intervals in p03_stats.py", "statistic": "number of run-level resamples"})
V.set("ci_level", LEVEL, "fraction", {"population": "bootstrap setting", "filter": "all intervals in p03_stats.py", "statistic": "coverage level of the percentile intervals"})
RNG = np.random.default_rng(SEED)
import seg_table as ST
D7_MIN_DUR, D7_TRIM, D7_PRE, D7_SPEED_TOL = 3.0, 1.0, (1.2, 0.2), 0.02
_E = lambda f, st: {"population": "analysis setting", "filter": f, "statistic": st}
V.set("d5_min_rpm", ST.MIN_SPINDLE_HZ_FOR_COMB * 60, "rpm", _E("D5 definition", "lowest spindle speed at which the harmonic comb leaves clean gaps"))
V.set("d5_band_lo_hz", ST.NH_BAND[0], "Hz", _E("D5 definition", "lower edge of the searched band"))
V.set("d5_band_hi_hz", ST.NH_BAND[1], "Hz", _E("D5 definition", "upper edge of the searched band"))
V.set("d5_comb_bins", ST.COMB_HALF_BINS, "bins", _E("D5 definition", "bins excluded on each side of a spindle harmonic"))
V.set("d7_min_dur_s", D7_MIN_DUR, "s", _E("D7 definition", "shortest path used"))
V.set("d7_trim_s", D7_TRIM, "s", _E("D7 definition", "time trimmed from the start of each path"))
V.set("d7_pre_from_s", D7_PRE[0], "s", _E("D7 definition", "start of the pre-cut window before the path"))
V.set("d7_pre_to_s", D7_PRE[1], "s", _E("D7 definition", "end of the pre-cut window before the path"))
V.set("d7_speed_tol", D7_SPEED_TOL, "fraction", _E("D7 definition", "tolerance on pre-cut spindle speed against N"))


def E(pop, filt, stat):
    return {"population": pop, "filter": filt, "statistic": stat}


class Boot:
    """Run-level resample counts, drawn once per population of runs (stratified by machine)."""
    cache = {}

    def __init__(self, df):
        runs = df[["machine", "run_id"]].drop_duplicates().sort_values(["machine", "run_id"])
        self.runs = list(runs.run_id)
        self.idx = {r: i for i, r in enumerate(self.runs)}
        key = tuple(self.runs)
        if key not in Boot.cache:
            W = np.zeros((B, len(self.runs)))
            for m, g in runs.groupby("machine"):
                cols = [self.idx[r] for r in g.run_id]
                draws = RNG.integers(0, len(cols), size=(B, len(cols)))
                for j in range(len(cols)):
                    np.add.at(W, (np.arange(B), np.array(cols)[draws[:, j]]), 1)
            Boot.cache[key] = W
        self.W = Boot.cache[key]

    def vec(self, df, values):
        """Sum of `values` per run, as a vector over self.runs."""
        s = pd.Series(values, index=df.index).groupby(df.run_id).sum()
        v = np.zeros(len(self.runs))
        for r, x in s.items():
            v[self.idx[r]] = x
        return v


def ci(arr):
    arr = np.asarray(arr, float); arr = arr[np.isfinite(arr)]
    a = (1 - LEVEL) / 2
    return float(np.quantile(arr, a)), float(np.quantile(arr, 1 - a)), int(len(arr))


def put(name, point, interval, unit, est, method):
    V.set(name, float(point), unit, est)
    lo, hi, nb = interval
    V.interval(name + "_ci", name, lo, hi, LEVEL, f"{method}; {nb} of {B} resamples usable", est)


def ratio_boot(bt, num_r, den_r):
    num = bt.W @ num_r; den = bt.W @ den_r
    with np.errstate(invalid="ignore", divide="ignore"):
        return num / den


def pooled_auc(df, col, pos=1, neg=0):
    """AUC over all label pos against label neg paths, with its bootstrap draws."""
    d = df[df.label.isin([pos, neg]) & df[col].notna()]
    bt = Boot(d)
    R = len(bt.runs)
    C = np.zeros((R, R)); n1 = np.zeros(R); n0 = np.zeros(R)
    g1 = {r: g[col].values for r, g in d[d.label == pos].groupby("run_id")}
    g0 = {r: np.sort(g[col].values) for r, g in d[d.label == neg].groupby("run_id")}
    for r, x in g1.items():
        n1[bt.idx[r]] = len(x)
        for s, y in g0.items():
            lt = np.searchsorted(y, x, side="left"); le = np.searchsorted(y, x, side="right")
            C[bt.idx[r], bt.idx[s]] = lt.sum() + 0.5 * (le - lt).sum()
    for s, y in g0.items():
        n0[bt.idx[s]] = len(y)
    point = C.sum() / (n1.sum() * n0.sum())
    num = np.einsum("br,rs,bs->b", bt.W, C, bt.W)
    den = (bt.W @ n1) * (bt.W @ n0)
    with np.errstate(invalid="ignore", divide="ignore"):
        draws = num / den
    return point, draws, int(n1.sum()), int(n0.sum())


def within_auc(df, col, strata):
    """Probability that a label 1 path exceeds a label 0 path of the same stratum (strata nested in runs)."""
    d = df[df.label.isin([0, 1]) & df[col].notna()]
    bt = Boot(d)
    num = np.zeros(len(bt.runs)); den = np.zeros(len(bt.runs))
    for k, g in d.groupby(strata):
        x1 = g[g.label == 1][col].values; x0 = g[g.label == 0][col].values
        if len(x1) and len(x0):
            r = bt.idx[g.run_id.iloc[0]]
            num[r] += (x1[:, None] > x0[None, :]).sum() + 0.5 * (x1[:, None] == x0[None, :]).sum()
            den[r] += len(x1) * len(x0)
    return num.sum() / den.sum(), ratio_boot(bt, num, den), int((den > 0).sum())


# ------------------------------------------------------------- D2 labels and cutting parameters --------------------
L = T[T.machine.isin(LAB)].copy()
# (a) up against down milling in design pairs (same run, tool, ap, ae, N, F)
pairs = []
for (rid, tool, ap, ae, N, F), g in L[L.label.isin([0, 1])].groupby(["run_id", "Tool", "ap", "ae", "N", "F"]):
    ups = g[g.direction == "UP"].sort_values("start"); dns = g[g.direction == "DOWN"].sort_values("start")
    up = ups.label.values; dn = dns.label.values
    for i in range(min(len(up), len(dn))):
        pairs.append({"machine": g.machine.iloc[0], "run_id": rid, "up": int(up[i]), "down": int(dn[i]),
                      "up_first": bool(ups.start.values[i] < dns.start.values[i])})
P = pd.DataFrame(pairs)
V.set("pairs_up_first_share", float(P.up_first.mean()), "fraction",
      E("design pairs of up and down paths", "laboratory machines, same run, tool, ap, ae, N and F, labels 0 or 1", "share of pairs in which the up path starts first"))
V.set("pairs_n", int(len(P)), "pairs", E("design pairs of up and down paths", "laboratory machines, same run, tool, ap, ae, N and F, labels 0 or 1", "count"))
CC = 0.5  # added to each discordant count so that a zero count still gives a finite odds ratio
V.set("d2_or_cc", CC, "pairs", E("analysis setting", "D2 odds ratios", "continuity correction added to each discordant count"))
for mach in LAB:
    d = P[P.machine == mach]
    bt = Boot(d)
    b_r = bt.vec(d, ((d.up == 1) & (d.down == 0)).astype(float).values)
    c_r = bt.vec(d, ((d.up == 0) & (d.down == 1)).astype(float).values)
    est = E("design pairs of up and down paths", f"{mach}, same run, tool, ap, ae, N and F, labels 0 or 1",
            "odds ratio of label 1 for up against down milling from discordant pairs, 0.5 added to each count")
    V.set(f"d2_pairs_{mach}", int(len(d)), "pairs", {**est, "statistic": "count of pairs"})
    V.set(f"d2_disc_up_{mach}", int(b_r.sum()), "pairs", {**est, "statistic": "pairs with up labelled 1 and down labelled 0"})
    V.set(f"d2_disc_down_{mach}", int(c_r.sum()), "pairs", {**est, "statistic": "pairs with down labelled 1 and up labelled 0"})
    V.set(f"d2_concord1_{mach}", int(((d.up == 1) & (d.down == 1)).sum()), "pairs", {**est, "statistic": "pairs with both labelled 1"})
    draws = (bt.W @ b_r + CC) / (bt.W @ c_r + CC)
    put(f"d2_or_updown_{mach}", (b_r.sum() + CC) / (c_r.sum() + CC), ci(draws), "odds ratio", est, "run bootstrap")

# (b) aluminium against steels, laboratory machines
L["alu"] = (L.Workpiece == "AA 6061").astype(float)
L["l1"] = (L.label == 1).astype(float)
bt = Boot(L)
a1 = bt.vec(L, (L.alu * L.l1).values); an = bt.vec(L, L.alu.values)
s1 = bt.vec(L, ((1 - L.alu) * L.l1).values); sn = bt.vec(L, (1 - L.alu).values)
est = E("labelled cutting paths", "laboratory machines, labels 0, 1 and 2", "share of paths with label 1 on AA 6061 minus the share on the two steels")
draws = (bt.W @ a1) / (bt.W @ an) - (bt.W @ s1) / (bt.W @ sn)
put("d2_alu_minus_steel", a1.sum() / an.sum() - s1.sum() / sn.sum(), ci(draws), "fraction", est, "run bootstrap")
V.set("d2_alu_runs", int(L[L.alu == 1].run_id.nunique()), "runs", E("laboratory runs", "AA 6061", "count"))
V.set("d2_steel_runs", int(L[L.alu == 0].run_id.nunique()), "runs", E("laboratory runs", "AISI 1018 or AISI 4140", "count"))

# (c) within-run association of label 1 with each cutting parameter
L["up"] = (L.direction == "UP").astype(float)
V.set("d2_within_runs", int(L[L.label.isin([0, 1])].groupby("run_id").label.nunique().eq(2).sum()), "runs",
      E("laboratory runs", "runs with both label 0 and label 1 paths", "count"))
for c in ("ap", "ae", "N", "fz", "mrr", "up"):
    p, draws, nr = within_auc(L, c, ["run_id"])
    est = E("laboratory paths, labels 0 and 1", "pairs of a label 1 and a label 0 path from the same run",
            f"probability that the label 1 path has the larger {c} (ties count half)")
    put(f"d2_within_auc_{c.lower()}", p, ci(draws), "AUC", est, "run bootstrap")

# ------------------------------------------------------------- D3 single-channel separation (AUC) --------------------
CHS = {"acc0": "acc0_rms", "acc1": "acc1_rms", "s0": "s0_rms", "s1": "s1_rms", "s2": "s2_rms",
       "load": "mtc_S_Axis_Load", "smarms": "mtc_smarms", "watt": "mtc_watt"}
d3 = []
for mach in LAB:
    sub = T[(T.machine == mach) & T.label.isin([0, 1])]
    for ch, col in CHS.items():
        if sub[col].notna().sum() == 0:
            continue
        p, draws, n1, n0 = pooled_auc(sub, col)
        est = E("labelled cutting paths", f"{mach}, labels 0 and 1, {col} available", f"AUC of {col} for label 1 against label 0")
        put(f"d3_auc_{mach}_{ch}", p, ci(draws), "AUC", est, "run bootstrap")
        d3.append({"machine": mach, "ch": ch, "auc": p, "ci": ci(draws)[:2], "n1": n1, "n0": n0})
        V.set(f"d3_n1_{mach}_{ch}", n1, "paths", {**est, "statistic": "label 1 paths"})
        V.set(f"d3_n0_{mach}_{ch}", n0, "paths", {**est, "statistic": "label 0 paths"})
res["d3"] = d3

# ------------------------------------------------------------- D4 matched strata ------------------------------------------
d4 = []
for mach in LAB:
    sub = T[(T.machine == mach) & T.label.isin([0, 1])]
    for ch, col in CHS.items():
        if sub[col].notna().sum() == 0:
            continue
        p, draws, nr = within_auc(sub, col, ["run_id", "Tool", "ap", "ae"])
        est = E("labelled cutting paths", f"{mach}, labels 0 and 1, strata of run, tool, ap and ae holding both labels", f"within-stratum AUC of {col}")
        put(f"d4_auc_{mach}_{ch}", p, ci(draws), "AUC", est, "run bootstrap")
        d4.append({"machine": mach, "ch": ch, "auc": p, "ci": ci(draws)[:2], "runs": nr})
    strata = sub.groupby(["run_id", "Tool", "ap", "ae"]).label.agg(lambda s: s.nunique() == 2)
    V.set(f"d4_strata_{mach}", int(strata.sum()), "strata", E("labelled cutting paths", f"{mach}, same run, tool, ap and ae", "strata holding both labels"))
    V.set(f"d4_runs_{mach}", int(strata[strata].reset_index().run_id.nunique()), "runs", E("labelled cutting paths", f"{mach}, same run, tool, ap and ae", "runs with at least one stratum holding both labels"))
res["d4"] = d4

# ------------------------------------------------------------- D5 non-harmonic dominant peak ------------------------
d5 = []
for mach in LAB:
    for ch in ("acc0", "acc1", "s1", "s2"):
        c1, c2 = f"{ch}_nh_peak_ratio", f"{ch}_h_peak_ratio"
        if c1 not in T or T[T.machine == mach][c1].notna().sum() == 0:
            continue
        sub = T[(T.machine == mach) & T.label.isin([0, 1]) & T[c1].notna()].copy()
        sub["nhdom"] = (sub[c1] > sub[c2]).astype(float)
        bt = Boot(sub)
        pieces = {}
        for lab in (0, 1):
            k = bt.vec(sub, (sub.nhdom * (sub.label == lab)).values); n = bt.vec(sub, (sub.label == lab).astype(float).values)
            pieces[lab] = (k, n)
            est = E("labelled cutting paths", f"{mach}, {ch}, spindle speed at least 3000 rpm, label {lab}",
                    "share of paths whose largest spectral line in 200 Hz to 10 kHz is not within two bins of a spindle harmonic")
            put(f"d5_nhdom_{mach}_{ch}_l{lab}", k.sum() / n.sum(), ci(ratio_boot(bt, k, n)), "fraction", est, "run bootstrap")
        (k1, n1), (k0, n0) = pieces[1], pieces[0]
        est = E("labelled cutting paths", f"{mach}, {ch}, spindle speed at least 3000 rpm, labels 0 and 1", "difference in non-harmonic dominant share, label 1 minus label 0")
        draws = ratio_boot(bt, k1, n1) - ratio_boot(bt, k0, n0)
        put(f"d5_nhdom_diff_{mach}_{ch}", k1.sum() / n1.sum() - k0.sum() / n0.sum(), ci(draws), "fraction", est, "run bootstrap")
        V.set(f"d5_n_{mach}_{ch}", int(len(sub)), "paths", E("labelled cutting paths", f"{mach}, {ch}, spindle speed at least 3000 rpm, labels 0 and 1", "count"))
        d5.append({"machine": mach, "ch": ch, "n": int(len(sub)), "l1": float(k1.sum() / n1.sum()), "l0": float(k0.sum() / n0.sum())})
res["d5"] = d5
# met along the way (not pre-specified): chatter-like lines in label 0 paths of tool A on AA 6061 (imi_vm20i)
sub = T[(T.machine == "imi_vm20i") & (T.label == 0) & (T.tool_base == "A") & (T.Workpiece == "AA 6061") & T.acc0_nh_peak_ratio.notna()].copy()
sub["nhdom"] = (sub.acc0_nh_peak_ratio > sub.acc0_h_peak_ratio).astype(float)
bt = Boot(sub)
k = bt.vec(sub, sub.nhdom.values); n = bt.vec(sub, np.ones(len(sub)))
est = E("labelled cutting paths", "imi_vm20i, tool A on AA 6061, label 0, spindle speed at least 3000 rpm, Acc0",
        "share of paths whose largest spectral line in 200 Hz to 10 kHz is not within two bins of a spindle harmonic")
put("d5x_nhdom_toola_alu_l0", k.sum() / n.sum(), ci(ratio_boot(bt, k, n)), "fraction", est, "run bootstrap")
V.set("d5x_n_toola_alu_l0", int(len(sub)), "paths", {**est, "statistic": "count"})
V.set("d5x_runs_toola_alu_l0", int(sub.run_id.nunique()), "runs", {**est, "statistic": "runs"})

# ------------------------------------------------------------- D7 specific cutting energy -------------------------------
rows = []
for (mach, run), g in T[T.machine.isin(LAB) & (T.label == 0)].groupby(["machine", "run"]):
    df = pd.read_csv(glob.glob(f"{RAW}/{mach}/dataset{run}/*_mtc.csv")[0], encoding="utf-8-sig", usecols=["time", "watt", "Spindle_Speed"])
    t = df.time.values; w = df.watt.values.astype(float); sp = df.Spindle_Speed.values.astype(float)
    for _, r in g.iterrows():
        if r.dur < D7_MIN_DUR:
            continue
        m = (t >= r.start + D7_TRIM) & (t <= r.end)          # skip the first second (MTConnect lags the sound)
        pm = (t >= r.start - D7_PRE[0]) & (t <= r.start - D7_PRE[1])
        if m.sum() < 3 or pm.sum() < 2 or abs(np.median(sp[pm]) - r.N) / r.N > D7_SPEED_TOL:
            continue
        rows.append({"machine": mach, "run_id": r.run_id, "Workpiece": r.Workpiece, "mrr": r.mrr, "p_cut": np.median(w[m]) - np.median(w[pm])})
D7 = pd.DataFrame(rows)
V.set("d7_paths", int(len(D7)), "paths", E("laboratory label 0 paths", "at least 3 s long, pre-cut spindle speed within 2 percent of N, first second trimmed", "count"))
d7 = []
for mach in LAB:
    for wp in ("AA 6061", "AISI 1018", "AISI 4140"):
        d = D7[(D7.machine == mach) & (D7.Workpiece == wp)]
        if len(d) < 10:
            continue
        key = f"{mach}_{wp.replace(' ', '').lower()}"
        bt = Boot(d)
        sxy = bt.vec(d, (d.p_cut * d.mrr).values); sxx = bt.vec(d, (d.mrr ** 2).values)
        est = E("laboratory label 0 paths", f"{mach}, {wp}, at least 3 s long, pre-cut spindle speed within 2 percent of N, first second trimmed",
                "least-squares slope through the origin of power above the pre-cut level (W) on programmed MRR (mm3/s), in J/mm3")
        put(f"d7_u_{key}", sxy.sum() / sxx.sum(), ci(ratio_boot(bt, sxy, sxx)), "J/mm3", est, "run bootstrap")
        V.set(f"d7_n_{key}", int(len(d)), "paths", {**est, "statistic": "count"})
        V.set(f"d7_runs_{key}", int(d.run_id.nunique()), "runs", {**est, "statistic": "runs"})
        V.set(f"d7_negative_share_{key}", float((d.p_cut <= 0).mean()), "fraction", {**est, "statistic": "share of paths with power at or below the pre-cut level"})
        d7.append({"machine": mach, "wp": wp, "u": sxy.sum() / sxx.sum(), "n": int(len(d)), "runs": int(d.run_id.nunique())})
res["d7"] = d7
D7.to_csv(f"{OUT}/d7_paths.csv", index=False)

# ------------------------------------------------------------- D10 defective tools against good counterparts -------
d10 = []
key = ["machine", "tool_base", "Workpiece", "ap", "ae", "N", "F", "direction"]
good = T[(~T.defective) & (T.label == 0)]
bad = T[(T.defective) & (T.label == 2)]
for ch in ("acc0", "acc1", "s0", "s1", "s2"):
    col = f"{ch}_rms"
    gm = good.groupby(key)[col].median().rename("good").reset_index()
    m = bad.merge(gm, on=key, how="inner")
    m = m[m[col].notna() & m.good.notna()].copy()
    if len(m) < 10:
        continue
    m["lr"] = np.log(m[col] / m.good)
    bt = Boot(m)
    runs = {r: g.lr.values for r, g in m.groupby("run_id")}
    draws = []
    for b in range(B):
        vals = np.concatenate([np.repeat(runs[r], int(bt.W[b, bt.idx[r]])) for r in bt.runs if bt.W[b, bt.idx[r]] > 0] or [np.array([np.nan])])
        draws.append(np.median(vals))
    est = E("label 2 paths with a label 0 path of the good counterpart tool on the same machine, workpiece, ap, ae, N, F and direction",
            f"{ch} RMS", "median log ratio of RMS, defective over matched good (reported as a ratio)")
    put(f"d10_ratio_{ch}", float(np.exp(np.median(m.lr))), ci(np.exp(np.array(draws))), "ratio", est, "run bootstrap of the median log ratio, exponentiated")
    V.set(f"d10_n_{ch}", int(len(m)), "paths", {**est, "statistic": "matched label 2 paths"})
    V.set(f"d10_runs_{ch}", int(m.run_id.nunique()), "runs", {**est, "statistic": "runs with matched label 2 paths"})
    d10.append({"ch": ch, "n": int(len(m)), "runs": int(m.run_id.nunique()), "ratio": float(np.exp(np.median(m.lr)))})
res["d10"] = d10

json.dump(res, open(f"{OUT}/stats.json", "w"), indent=1, default=float)
V.save()
print("p03_stats: registered", len(V._touched), "values")