"""p01_inventory.py: what was received, and does it match the documentation?

Inputs: raw/ (small files staged from the device), raw/Phase2/derived/* (device-side passes
d01 and d02 over the accelerometer CSVs, WAV and MP4 files, and the full file inventory),
inputs/sources.json. Outputs: outputs/inventory.json, outputs/inventory_runs.csv and values.
"""
import json, glob, os, re, sys
import numpy as np, pandas as pd
sys.path.insert(0, "/home/claude/mill/code")
from vals import Values

RAW = "/home/claude/mill/raw"
DER = f"{RAW}/Phase2/derived"
OUT = "/home/claude/mill/outputs"
MACH = ("imi_vm20i", "imi_vmx30ui", "tmf_vf10")
V = Values(script="p01_inventory.py")
S = json.load(open("/home/claude/mill/inputs/sources.json"))
src = {s["id"]: s for s in S["sources"]}

REL = {"population": "Kaggle release as received (folder archive (1)), Phase2 outputs excluded",
       "filter": "none", "statistic": "count"}

# ---------------- files ----------------------------------------------------------------------------
inv = json.load(open(f"{DER}/file_inventory.json"))
inv = [r for r in inv if not r["path"].startswith("Phase2/")]
V.set("n_files", len(inv), "files", REL)
V.set("bytes_total", sum(r["bytes"] for r in inv), "bytes", {**REL, "statistic": "sum of file sizes"})
V.set("bytes_match_kaggle", int(sum(r["bytes"] for r in inv) == V.get("kaggle_total_bytes") if "kaggle_total_bytes" in V.v else False),
      "boolean", {**REL, "statistic": "1 if the byte total equals Kaggle totalBytes"})


def kind(p):
    b = p.split("/")[-1]
    if b.endswith("_acc.csv"): return "acc"
    if b.endswith("_mtc.csv"): return "mtc"
    if b.endswith("_ct.csv"): return "ct"
    if b == "label.csv": return "label"
    m = re.search(r"_(s\d)\.wav$", b)
    if m: return m.group(1)
    if b.endswith(".mp4"): return "video"
    if b.endswith(".nc"): return "nc"
    if b.endswith(".png"): return "png"
    if b.endswith(".xml"): return "xml"
    if b.endswith(".xlsx"): return "xlsx"
    return "other"


for r in inv:
    r["kind"] = kind(r["path"])
kinds = pd.Series([r["kind"] for r in inv]).value_counts().to_dict()
for k, n in kinds.items():
    V.set(f"n_files_{k}", int(n), "files", {**REL, "filter": f"file kind {k}"})
V.set("bytes_acc", sum(r["bytes"] for r in inv if r["kind"] == "acc"), "bytes", {**REL, "filter": "acc.csv files", "statistic": "sum of sizes"})
V.set("bytes_wav", sum(r["bytes"] for r in inv if r["kind"] in ("s0", "s1", "s2")), "bytes", {**REL, "filter": "WAV files", "statistic": "sum of sizes"})
V.set("acc_file_max_bytes", max(r["bytes"] for r in inv if r["kind"] == "acc"), "bytes", {**REL, "filter": "acc.csv files", "statistic": "max size"})

# exact duplicate small files (same sha256 in two places) other than the NC programs reused across runs
sha = pd.DataFrame([r for r in inv if "sha256" in r])
dup = sha[sha.duplicated("sha256", keep=False)]
dup_non_nc = dup[~dup.kind.isin(["nc"])]
V.set("n_dup_small_files_non_nc", int(len(dup_non_nc)), "files",
      {**REL, "filter": "files under 50 MB other than NC programs", "statistic": "files whose sha256 equals another file's"})
V.set("n_dup_nc_files", int(len(dup[dup.kind == "nc"])), "files",
      {**REL, "filter": "NC programs", "statistic": "files whose sha256 equals another NC file"})

# ---------------- runs -------------------------------------------------------------------------------
summ = pd.read_excel(f"{RAW}/dataset_summary.xlsx")
summ["run"] = summ["dataset"].astype(int)
runs = []
for mach in MACH:
    for d in sorted(glob.glob(f"{RAW}/{mach}/dataset*"), key=lambda p: int(p.split("dataset")[-1])):
        runs.append((mach, int(d.split("dataset")[-1])))
V.set("n_runs", len(runs), "runs", {**REL, "statistic": "count of dataset folders"})
for mach in MACH:
    V.set(f"n_runs_{mach}", sum(1 for m, _ in runs if m == mach), "runs", {**REL, "filter": mach, "statistic": "count of dataset folders"})
V.set("n_summary_rows", len(summ), "rows", {"population": "dataset_summary.xlsx", "filter": "none", "statistic": "rows"})
V.set("summary_matches_folders", int(set(zip(summ.machine, summ.run)) == set(runs)), "boolean",
      {"population": "dataset_summary.xlsx vs dataset folders", "filter": "none", "statistic": "1 if the (machine, dataset) sets are equal"})

acc_stats = {os.path.basename(f)[:-5]: json.load(open(f)) for f in glob.glob(f"{DER}/acc_stats/*.json")}
aud_stats = {os.path.basename(f)[:-5]: json.load(open(f)) for f in glob.glob(f"{DER}/audio_stats/*.json")}
vid = {(v["machine"], int(v["run"].replace("dataset", ""))): v for v in json.load(open(f"{DER}/video_probe.json"))}

rows = []
for mach, run in runs:
    d = f"{RAW}/{mach}/dataset{run}"
    files = [r for r in inv if r["path"].startswith(f"{mach}/dataset{run}/")]
    ks = {r["kind"] for r in files}
    lab = pd.read_csv(f"{d}/label.csv")
    mtc_f = glob.glob(f"{d}/*_mtc.csv")[0]
    mtc = pd.read_csv(mtc_f, encoding="utf-8-sig", usecols=["time"])
    stems = {re.sub(r"_(acc|mtc|ct|s\d)\.(csv|wav)$", "", os.path.basename(r["path"])) for r in files if r["kind"] in ("acc", "mtc", "ct", "s0", "s1", "s2")}
    rec = {"machine": mach, "run": run, "n_files": len(files), "kinds": sorted(ks), "n_stems": len(stems), "stem": sorted(stems)[0],
           "label_rows": len(lab), "label_first_start": float(lab.start.min()), "label_last_end": float(lab.end.max()),
           "cut_s": float((lab.end - lab.start).sum()),
           "mtc_rows": len(mtc), "mtc_t_last": float(mtc.time.iloc[-1]),
           "label_cols": int(lab.shape[1]), "mtc_cols": len(pd.read_csv(mtc_f, encoding="utf-8-sig", nrows=0).columns)}
    tag = f"{mach}_dataset{run}"
    if tag in acc_stats:
        a = acc_stats[tag]; rec.update({"acc_duration_s": a["duration_s"], "acc_rows": a["rows"], "acc_fs": a["fs_est"],
                                        "acc_dt_off": a["n_dt_off_1e-9"], "acc_nan": a["nan_acc0"] + a["nan_acc1"] + a["nan_time"]})
    for ch in ("s0", "s1", "s2"):
        k = f"{tag}_{ch}"
        if k in aud_stats:
            w = aud_stats[k]
            rec[f"{ch}_duration_s"] = w["duration_s"]; rec[f"{ch}_fs"] = w["framerate"]; rec[f"{ch}_bits"] = 8 * w["sampwidth"]
            rec[f"{ch}_nch"] = w["nchannels"]; rec[f"{ch}_clip"] = w["n_fullscale_pos"] + w["n_fullscale_neg"]
    ctf = glob.glob(f"{d}/*_ct.csv")
    if ctf:
        ct = pd.read_csv(ctf[0]); rec["ct_rows"] = len(ct); rec["ct_t_last"] = float(ct.time.iloc[-1]); rec["ct_cols"] = int(ct.shape[1])
        rec["ct_dt_median"] = float(np.median(np.diff(ct.time.values)))
    if (mach, run) in vid:
        v = vid[(mach, run)]; rec["video_duration_s"] = v["duration_s"]
        vs = re.match(r"(\d{8})_(\d{6})(\d+)", v["file"].split("_comp")[0].replace(".", ""))
        # reference stem, e.g. 20240229_053831.227000; video stem, e.g. 20240229_053833440 or 20240229_044455.365
        ref = pd.to_datetime(rec["stem"], format="%Y%m%d_%H%M%S.%f")
        vstem = v["file"].split("_comp")[0]
        if "." in vstem:
            vt = pd.to_datetime(vstem, format="%Y%m%d_%H%M%S.%f")
        else:
            vt = pd.to_datetime(vstem[:15], format="%Y%m%d_%H%M%S") + pd.to_timedelta(int(vstem[15:]), unit="ms")
        rec["video_start_offset_s"] = (vt - ref).total_seconds()
    ncs = [r for r in files if r["kind"] == "nc"]
    rec["nc_file"] = os.path.basename(ncs[0]["path"]) if ncs else None
    srow = summ[(summ.machine == mach) & (summ.run == run)].iloc[0]
    rec["summary_tool"] = str(srow["Tool"]); rec["summary_workpiece"] = str(srow["Workpiece"])
    rec["summary_remark"] = None if pd.isna(srow["Remark"]) else str(srow["Remark"]).strip()
    rec["label_tools"] = sorted(lab.Tool.astype(str).unique()); rec["label_workpieces"] = sorted(lab.Workpiece.astype(str).unique())
    rows.append(rec)
R = pd.DataFrame(rows)
os.makedirs(OUT, exist_ok=True)
R.to_csv(f"{OUT}/inventory_runs.csv", index=False)

# ---------------- shapes (columns and rows per file type) -----------------------------------------
SH = {"population": "Kaggle release as received (folder archive (1)), Phase2 outputs excluded"}
for col, lab_ in (("label_cols", "label.csv"), ("ct_cols", "ct.csv")):
    x = R[col].dropna()
    V.set(f"{col}_min", int(x.min()), "columns", {**SH, "filter": f"{lab_} files ({len(x)})", "statistic": "fewest columns"})
    V.set(f"{col}_max", int(x.max()), "columns", {**SH, "filter": f"{lab_} files ({len(x)})", "statistic": "most columns"})
for mach in MACH:
    x = R[R.machine == mach].mtc_cols
    V.set(f"mtc_cols_{mach}_min", int(x.min()), "columns", {**SH, "filter": f"{mach} mtc.csv files ({len(x)})", "statistic": "fewest columns"})
    V.set(f"mtc_cols_{mach}_max", int(x.max()), "columns", {**SH, "filter": f"{mach} mtc.csv files ({len(x)})", "statistic": "most columns"})
V.set("ct_rows_total", int(R.ct_rows.dropna().sum()), "rows", {**SH, "filter": "ct.csv files", "statistic": "rows"})
acc_cols = [len(eval(a["columns"])) if isinstance(a["columns"], str) else len(a["columns"]) for a in acc_stats.values()]
V.set("acc_cols_min", int(min(acc_cols)), "columns", {**SH, "filter": f"acc.csv files ({len(acc_cols)})", "statistic": "fewest columns"})
V.set("acc_cols_max", int(max(acc_cols)), "columns", {**SH, "filter": f"acc.csv files ({len(acc_cols)})", "statistic": "most columns"})

# ---------------- modality availability ----------------------------------------------------------
avail = {}
for mach in MACH:
    sub = R[R.machine == mach]
    for k in ("acc", "s0", "s1", "s2", "mtc", "ct", "video", "nc", "label"):
        n = int(sub.kinds.apply(lambda ks: k in ks).sum())
        avail[(mach, k)] = n
        V.set(f"n_runs_with_{k}_{mach}", n, "runs", {**REL, "filter": f"{mach} runs containing a {k} file", "statistic": "count"})
V.set("n_runs_one_stem", int((R.n_stems == 1).sum()), "runs",
      {**REL, "statistic": "runs whose sensor files share one reference timestamp stem"})

# documented availability (article): acc and Sound1 on imi_vm20i only; video and NC for lab runs;
# CT for tmf_vf10 only; Sound0, Sound2 and MTConnect on all machines
expect = {("imi_vm20i", "acc"): "all", ("imi_vmx30ui", "acc"): "none", ("tmf_vf10", "acc"): "none",
          ("imi_vm20i", "s1"): "all", ("imi_vmx30ui", "s1"): "none", ("tmf_vf10", "s1"): "none",
          ("imi_vm20i", "video"): "all", ("imi_vmx30ui", "video"): "all", ("tmf_vf10", "video"): "none",
          ("imi_vm20i", "nc"): "all", ("imi_vmx30ui", "nc"): "all", ("tmf_vf10", "nc"): "none",
          ("imi_vm20i", "ct"): "none", ("imi_vmx30ui", "ct"): "none", ("tmf_vf10", "ct"): "all"}
for mach in MACH:
    for k in ("s0", "s2", "mtc", "label"):
        expect[(mach, k)] = "all"
mismatch = []
for (mach, k), e in expect.items():
    n_all = V.get(f"n_runs_{mach}")
    ok = (avail[(mach, k)] == n_all) if e == "all" else (avail[(mach, k)] == 0)
    if not ok:
        mismatch.append(f"{mach}:{k}")
V.set("n_availability_checks", len(expect), "checks", {"population": "machine by modality cells", "filter": "documented availability", "statistic": "count"})
V.set("n_availability_mismatch", len(mismatch), "checks", {"population": "machine by modality cells", "filter": "documented availability", "statistic": "cells where files disagree with the article"})

# ---------------- durations ------------------------------------------------------------------------------
R["dur_s0"] = R["s0_duration_s"]
for mach in MACH:
    sub = R[R.machine == mach]
    V.set(f"dur_s0_{mach}_s", float(sub.dur_s0.sum()), "s", {**REL, "filter": f"{mach}", "statistic": "sum over runs of Sound0 WAV duration"})
    V.set(f"dur_mtc_{mach}_s", float(sub.mtc_t_last.sum()), "s", {**REL, "filter": f"{mach}", "statistic": "sum over runs of the last MTConnect time stamp"})
    V.set(f"cut_{mach}_s", float(sub.cut_s.sum()), "s", {**REL, "filter": f"{mach} label rows", "statistic": "sum of end minus start"})
V.set("dur_s0_total_s", float(R.dur_s0.sum()), "s", {**REL, "statistic": "sum over runs of Sound0 WAV duration"})
V.set("cut_total_s", float(R.cut_s.sum()), "s", {**REL, "filter": "all label rows", "statistic": "sum of end minus start"})
V.set("cut_share_s0", float(R.cut_s.sum() / R.dur_s0.sum()), "fraction", {**REL, "filter": "all runs", "statistic": "labelled cutting time over Sound0 duration"})
# which recorded duration reproduces the article's per-machine durations (to the second)?
cands = {"s0": "dur_s0", "mtc_last": "mtc_t_last"}
match = {}
for c, col in cands.items():
    match[c] = all(abs(round(R[R.machine == m][col].sum()) - V.get(f"doc_duration_{m}_s")) <= 1 for m in MACH)
V.set("doc_durations_match_s0", int(match["s0"]), "boolean", {**REL, "statistic": "1 if Sound0 totals equal the article's per-machine durations within 1 s"})
V.set("doc_durations_match_mtc", int(match["mtc_last"]), "boolean", {**REL, "statistic": "1 if MTConnect span totals equal the article's per-machine durations within 1 s"})
V.set("doc_cut_time_match", int(abs(round(R.cut_s.sum()) - V.get("doc_cut_time_s")) <= 1), "boolean",
      {**REL, "statistic": "1 if total labelled cutting time equals the article's cutting time within 1 s"})
V.set("dur_s0_minus_doc_maxabs_s", float(max(abs(R[R.machine == m].dur_s0.sum() - V.get(f"doc_duration_{m}_s")) for m in MACH)), "s",
      {**REL, "filter": "per machine", "statistic": "max absolute difference between the Sound0 total and the article's printed duration"})
V.set("cut_minus_doc_abs_s", float(abs(R.cut_s.sum() - V.get("doc_cut_time_s"))), "s",
      {**REL, "filter": "all label rows", "statistic": "absolute difference between labelled cutting time and the article's printed cutting time"})

# per-run agreement between modality durations (s)
lab_runs = R[R.machine != "tmf_vf10"]
R["d_s2_minus_s0"] = R["s2_duration_s"] - R["s0_duration_s"]
R["d_mtc_minus_s0"] = R["mtc_t_last"] - R["s0_duration_s"]
R["d_acc_minus_s0"] = R.get("acc_duration_s") - R["s0_duration_s"]
R["d_s1_minus_s0"] = R.get("s1_duration_s") - R["s0_duration_s"]
R["d_label_end_minus_s0"] = R["label_last_end"] - R["s0_duration_s"]
for col in ("d_s2_minus_s0", "d_mtc_minus_s0", "d_acc_minus_s0", "d_s1_minus_s0", "d_label_end_minus_s0"):
    x = R[col].dropna()
    V.set(f"{col}_maxabs", float(x.abs().max()), "s", {**REL, "filter": f"runs with both streams ({len(x)})", "statistic": "max absolute difference in duration"})
    V.set(f"{col}_median", float(x.median()), "s", {**REL, "filter": f"runs with both streams ({len(x)})", "statistic": "median difference in duration"})
V.set("n_label_end_beyond_s0", int((R.d_label_end_minus_s0 > 0).sum()), "runs", {**REL, "statistic": "runs whose last label end is after the end of Sound0"})

# ---------------- sampling ------------------------------------------------------------------------------
fs_acc = np.array([a["fs_est"] for a in acc_stats.values()])
V.set("acc_fs_est_min", float(fs_acc.min()), "Hz", {"population": "acc.csv files", "filter": "all", "statistic": "min of 1/median time step"})
V.set("acc_fs_est_max", float(fs_acc.max()), "Hz", {"population": "acc.csv files", "filter": "all", "statistic": "max of 1/median time step"})
V.set("acc_n_dt_off_total", int(sum(a["n_dt_off_1e-9"] for a in acc_stats.values())), "steps",
      {"population": "acc.csv files", "filter": "all", "statistic": "time steps differing from the file median by more than 1e-9 s"})
V.set("acc_rows_total", int(sum(a["rows"] for a in acc_stats.values())), "rows", {"population": "acc.csv files", "filter": "all", "statistic": "rows"})
V.set("acc_nan_total", int(sum(a["nan_acc0"] + a["nan_acc1"] + a["nan_time"] for a in acc_stats.values())), "cells",
      {"population": "acc.csv files", "filter": "all", "statistic": "missing cells"})
W = pd.DataFrame(aud_stats.values())
V.set("wav_n", len(W), "files", {"population": "WAV files", "filter": "all", "statistic": "count"})
V.set("wav_all_48k_16bit_mono", int(((W.framerate == 48000) & (W.sampwidth == 2) & (W.nchannels == 1) & (W.comptype == "NONE")).all()), "boolean",
      {"population": "WAV files", "filter": "all", "statistic": "1 if every header is 48 kHz, 16-bit, mono, uncompressed"})
W["clip_n"] = W.n_fullscale_pos + W.n_fullscale_neg
V.set("wav_n_files_with_clipping", int((W.clip_n > 0).sum()), "files", {"population": "WAV files", "filter": "all", "statistic": "files with at least one full-scale sample"})
V.set("wav_clip_samples_total", int(W.clip_n.sum()), "samples", {"population": "WAV files", "filter": "all", "statistic": "full-scale samples"})
V.set("wav_samples_total", int(W.nframes.sum()), "samples", {"population": "WAV files", "filter": "all", "statistic": "samples"})

# MTConnect time steps
mt = {}
for mach in MACH:
    dts = []
    for f in sorted(glob.glob(f"{RAW}/{mach}/dataset*/*_mtc.csv")):
        t = pd.read_csv(f, encoding="utf-8-sig", usecols=["time"]).time.values
        dts.append(np.diff(t))
    d = np.concatenate(dts)
    E = {"population": f"{mach} mtc.csv rows", "filter": "consecutive rows within a run", "statistic": None}
    V.set(f"mtc_dt_median_{mach}", float(np.median(d)), "s", {**E, "statistic": "median time step"})
    V.set(f"mtc_dt_p99_{mach}", float(np.quantile(d, 0.99)), "s", {**E, "statistic": "99th percentile time step"})
    V.set(f"mtc_dt_max_{mach}", float(d.max()), "s", {**E, "statistic": "max time step"})
    V.set(f"mtc_share_dt_gt1_{mach}", float((d > 1.0).mean()), "fraction", {**E, "statistic": "share of steps over 1 s"})
    V.set(f"mtc_rows_{mach}", int(len(d) + len(dts)), "rows", {**E, "filter": "all", "statistic": "rows"})
ctd = R[R.machine == "tmf_vf10"].ct_dt_median
V.set("ct_dt_median_min", float(ctd.min()), "s", {"population": "ct.csv files", "filter": "tmf_vf10", "statistic": "min over runs of the median time step"})
V.set("ct_dt_median_max", float(ctd.max()), "s", {"population": "ct.csv files", "filter": "tmf_vf10", "statistic": "max over runs of the median time step"})

# video offsets
vo = R.video_start_offset_s.dropna()
V.set("video_offset_min_s", float(vo.min()), "s", {"population": "lab runs with video", "filter": "all", "statistic": "min of video file-name time minus reference stem time"})
V.set("video_offset_max_s", float(vo.max()), "s", {"population": "lab runs with video", "filter": "all", "statistic": "max of video file-name time minus reference stem time"})
V.set("video_offset_median_s", float(vo.median()), "s", {"population": "lab runs with video", "filter": "all", "statistic": "median of video file-name time minus reference stem time"})
V.set("n_video_negative_offset", int((vo < 0).sum()), "runs", {"population": "lab runs with video", "filter": "all", "statistic": "runs whose video starts before the reference stem"})

# ---------------- labels vs summary and tool table --------------------------------------------
bad_tool = R[[not (set(t.replace(" ", "").split(",")) == set(lt)) for t, lt in zip(R.summary_tool, R.label_tools)]]
bad_wp = R[[not ({w} == set(lw)) for w, lw in zip(R.summary_workpiece, R.label_workpieces)]]
V.set("n_runs_tool_mismatch_summary", len(bad_tool), "runs", {"population": "all runs", "filter": "none", "statistic": "runs whose label.csv tools differ from dataset_summary.xlsx"})
V.set("n_runs_wp_mismatch_summary", len(bad_wp), "runs", {"population": "all runs", "filter": "none", "statistic": "runs whose label.csv workpiece differs from dataset_summary.xlsx"})
V.set("n_runs_with_remark", int(R.summary_remark.notna().sum()), "runs", {"population": "dataset_summary.xlsx", "filter": "Remark not empty", "statistic": "count"})
V.set("n_runs_early_stop", int(R.summary_remark.fillna("").str.contains("Early stopped").sum()), "runs",
      {"population": "dataset_summary.xlsx", "filter": "Remark contains 'Early stopped'", "statistic": "count"})

# structure counts and envelope resolution used by the synchronisation checks
V.set("n_machines", int(len(MACH)), "machines", {**REL, "statistic": "count of machine folders"})
V.set("n_materials", int(summ.Workpiece.nunique()), "materials", {"population": "dataset_summary.xlsx", "filter": "none", "statistic": "distinct workpiece materials"})
env_steps = {round(a["env_block_samples"] * a["dt_median"], 6) for a in acc_stats.values()} | {1 / 10 for _ in [0]}
assert len(env_steps) == 1, env_steps
V.set("env_step_s", float(env_steps.pop()), "s", {"population": "10 Hz block-RMS envelopes of acc.csv and WAV files", "filter": "d01 and d02 device passes", "statistic": "block length"})
V.set("env_rate_hz", 10, "Hz", {"population": "10 Hz block-RMS envelopes of acc.csv and WAV files", "filter": "d01 and d02 device passes", "statistic": "envelope rate"})
inventory = {"runs": rows, "kinds": kinds, "availability": {f"{m}:{k}": n for (m, k), n in avail.items()},
             "availability_mismatch": mismatch, "tool_mismatch_runs": bad_tool[["machine", "run", "summary_tool", "label_tools"]].to_dict("records"),
             "wp_mismatch_runs": bad_wp[["machine", "run", "summary_workpiece", "label_workpieces"]].to_dict("records"),
             "dup_small_files_non_nc": dup_non_nc[["path", "sha256"]].to_dict("records")}
json.dump(inventory, open(f"{OUT}/inventory.json", "w"), indent=1, default=str)
V.save()
print("p01_inventory: registered", len(V._touched), "values; availability mismatches:", mismatch,
      "; tool mismatches:", len(bad_tool), "; wp mismatches:", len(bad_wp), "; dup non-nc:", len(dup_non_nc))