"""r09_esm.py: the Online Resources that are data tables (ESM_1.csv, ESM_3.csv, ESM_4.csv).

ESM_1: the label 0 and label 1 paths of imi_vm20i (label 2 is outside every analysis) with their fold in every
       evaluation scheme: the family and material folds, the run-held-out fold of each repeat and the random
       path-level fold of each repeat, regenerated with the seeds and random streams of the analysis. The first
       repeat's run-held-out out-of-fold scores are recomputed from these folds and asserted equal to the
       stored ones.
ESM_3: the off-harmonic indicators of the same paths (D5 settings) for Acc0, Acc1, Sound1 and Sound2, from
       r03_explore.py, and the frequency of the dominant off-harmonic line of Acc0 and Sound2 where the
       indicator is on (r11_review.py, deviation DV17).
ESM_4: the planned tables of exploratory analyses E1 and E5 (deviation DV15): the share of paths with more than
       d9_threshold of their samples at full scale by machine, sound channel and label, and the off-harmonic
       shares by channel, indicator setting and label with their run-bootstrap intervals, all read from the
       values files.
Each file starts with comment lines (#) carrying the article title, journal, authors and the corresponding
author's affiliation and e-mail address, as the guidelines ask for every supplementary file.
"""
import sys, json, re
sys.path.insert(0, "/home/claude/mill/phase4/code")
sys.path.insert(0, "/home/claude/mill/phase3/code")
sys.path.insert(0, "/home/claude/mill/code")
import r_freeze
r_freeze.assert_freeze()

import numpy as np
import pandas as pd
import q_common as QC

P4 = "/home/claude/mill/phase4"
OUT = f"{P4}/esm"
V2 = json.load(open("/home/claude/mill/outputs/values.json"))
V3 = json.load(open("/home/claude/mill/phase3/outputs/values3.json"))
g = lambda n: (V3[n] if n in V3 else V2[n])["value"]
REP, K, SEED, C = int(g("s1_repeats")), int(g("s1_folds")), int(g("seed_analysis")), float(g("lr_c"))


def header(num, caption, columns, ext="csv"):
    front = open(f"{P4}/tex/src/01_front.tex").read()
    title = re.search(r"\\title\[[^\]]*\]\{([^}]*)\}", front).group(1)
    A = json.load(open(f"{P4}/inputs/authors.json"))
    S4 = json.load(open(f"{P4}/inputs/sources4.json"))
    ca = next(a for a in A["authors"] if a.get("corresponding"))
    inst = S4["institutions"][ca["inst"]]
    names = ", ".join(f"{a['given']} {a['family']}" for a in A["authors"])
    lines = [f"Article: {title}", "Journal: The International Journal of Advanced Manufacturing Technology", f"Authors: {names}",
             f"Corresponding author: {ca['given']} {ca['family']}, {inst['name']}, {inst['city']}, {inst['state']}, {inst['country']}; {ca['email']}",
             f"Online Resource {num} (ESM_{num}.{ext}): {caption}", "Columns: " + "; ".join(f"{k}: {v}" for k, v in columns.items())]
    return "".join(f"# {l}\n" for l in lines)


import os
os.makedirs(OUT, exist_ok=True)
P = QC.load_paths()
d = pd.DataFrame({"machine": P.machine, "run": P.run.astype(int), "start_s": P.start.round(3), "end_s": P.end.round(3), "label": P.label.astype(int),
                  "tool": P.Tool, "tool_family": P.family, "workpiece": P.Workpiece, "spindle_speed_rpm": P.N.astype(int),
                  "fold_family": P.family, "fold_material": P.Workpiece})
rng = np.random.default_rng(SEED)
runf = []
for rep in range(REP):
    f = QC.run_folds(P.run_id.values, K, rng)
    d[f"run_fold_rep{rep + 1:02d}"] = f
    runf.append(f)
rng = np.random.default_rng(SEED)
idx = np.arange(len(P)); y = P.y.values
for rep in range(REP):
    folds = np.zeros(len(P), int)
    for lab in (0, 1):
        ii = idx[y == lab].copy(); rng.shuffle(ii)
        folds[ii] = np.arange(len(ii)) % K
    d[f"random_fold_rep{rep + 1:02d}"] = folds
# the first repeat's run-held-out scores, recomputed from the exported folds, must equal the stored scores
arr = np.load(f"{P4}/outputs/r01_arrays.npz", allow_pickle=True)
X = QC.design(P, "highrate"); s = np.full(len(P), np.nan)
for f in range(K):
    tr, te = runf[0] != f, runf[0] == f
    s[te] = QC.fit_predict(X[tr], y[tr], X[te], "lr", C, False, True)
assert np.allclose(s, arr["h1_scores_run"][0], rtol=0, atol=1e-12), "exported run folds do not reproduce the stored scores"
cols1 = {"machine, run, start_s, end_s": "the path as in the run's label file", "label": "0 normal, 1 abnormal", "tool, tool_family, workpiece, spindle_speed_rpm": "as recorded",
         "fold_family, fold_material": "the held-out group in the family and material schemes",
         "run_fold_repNN": f"fold (0 to {K - 1}) of the path's run in repeat NN of the run-held-out scheme",
         "random_fold_repNN": f"fold (0 to {K - 1}) of the path in repeat NN of the random path-level scheme (stratified by label)"}
with open(f"{OUT}/ESM_1.csv", "w") as fh:
    fh.write(header(1, f"fold assignments of the {len(d):,} label 0 and label 1 paths of imi_vm20i in every evaluation scheme", cols1))
    d.to_csv(fh, index=False)
ind = pd.read_csv(f"{P4}/outputs/per_path_indicators.csv")
assert len(ind) == len(P) and (ind.start.values == P.start.values).all()
lf = pd.read_csv(f"{P4}/outputs/line_frequencies.csv")
assert len(lf) == len(P) and (lf.start.values == P.start.values).all() and (lf.run.values == P.run.values).all()
e3 = pd.DataFrame({"machine": ind.machine, "run": ind.run.astype(int), "start_s": ind.start.round(3), "end_s": ind.end.round(3), "label": ind.label.astype(int),
                   "tool": ind.Tool, "tool_family": ind.family, "workpiece": ind.Workpiece, "spindle_speed_rpm": ind.N.astype(int), "duration_s": ind.dur.round(3)})
for ch, name in (("acc0", "acc0"), ("acc1", "acc1"), ("s1", "sound1"), ("s2", "sound2")):
    e3[f"offharmonic_{name}"] = ind[f"{ch}_offharmonic"].map(lambda v: "" if pd.isna(v) else int(v))
for ch, name in (("acc0", "acc0"), ("s2", "sound2")):
    on = e3[f"offharmonic_{name}"].astype(str) == "1"
    assert (on == lf[f"{ch}_line_hz"].notna()).all(), name       # a frequency exactly where the indicator is on
    e3[f"line_hz_{name}"] = lf[f"{ch}_line_hz"].map(lambda v: "" if pd.isna(v) else f"{v:.2f}")
cols3 = {"machine, run, start_s, end_s": "the path", "label": "0 normal, 1 abnormal", "duration_s": "end minus start",
         "offharmonic_*": f"1 when the largest spectral line between {g('d5_band_lo_hz'):.0f} Hz and {g('d5_band_hi_hz'):.0f} Hz lies more than {g('d5_comb_bins'):.0f} bins from every harmonic of the commanded spindle frequency, 0 otherwise, empty below {g('d5_min_rpm'):.0f} r/min (values from values.json)",
         "line_hz_*": "frequency in Hz of that largest line where the indicator is 1, empty otherwise"}
with open(f"{OUT}/ESM_3.csv", "w") as fh:
    fh.write(header(3, f"off-harmonic spectral indicators of the {len(e3):,} label 0 and label 1 paths of imi_vm20i, with the frequency of the dominant off-harmonic line of Acc0 and Sound2", cols3))
    e3.to_csv(fh, index=False)

# ------------------------------------------------------------------ ESM_4: the E1 and E5 tables, read from the values files
V4 = json.load(open(f"{P4}/outputs/values4.json"))
VA = {**V2, **V3, **V4}
CH = {"s0": "Sound0", "s1": "Sound1", "s2": "Sound2", "acc0": "Acc0", "acc1": "Acc1"}
rows = []
for mach in ("imi_vm20i", "imi_vmx30ui"):
    for ch in ("s0", "s1", "s2"):
        for lab in (0, 1):
            nm = f"e1_{mach}_{ch}_l{lab}_gt1pct"
            nm = nm if nm in VA else (f"d9_paths_gt1pct_{mach}_{ch}_l{lab}" if f"d9_paths_gt1pct_{mach}_{ch}_l{lab}" in VA else None)
            if nm is None:
                continue
            rows.append({"analysis": "E1", "machine": mach, "channel": CH[ch], "setting": f"more than {100 * g('d9_threshold'):.0f} percent of samples at full scale",
                         "label": lab, "paths": VA[f"e1_{mach}_{ch}_l{lab}_n"]["value"], "share": round(VA[nm]["value"], 4), "ci_lo": "", "ci_hi": "", "value_name": nm})
SETN = {"base": "D5 settings (comb of {} bins, band {:.0f} Hz to {:.0f} Hz)".format(int(g("d5_comb_bins")), g("d5_band_lo_hz"), g("d5_band_hi_hz")),
        "comb1": "narrow comb of {} bin (V6a)".format(int(g("comb_bins_v6a"))), "comb3": "wide comb of {} bins (V6b)".format(int(g("comb_bins_v6b"))),
        "rel": "relative comb, {:.1f} percent of each harmonic's frequency (V6c)".format(100 * g("rel_speed_tol")),
        "band5k": "band up to {:.0f} Hz (V7)".format(g("band_hi_v7_hz"))}
yv = P.y.values
A3 = np.load(f"{P4}/outputs/r03_arrays.npz")
for ch in ("acc0", "acc1", "s1", "s2"):
    for sname in ("base", "comb1", "comb3", "rel", "band5k"):
        for lab in (0, 1):
            nm = f"e5_{ch}_{sname}_l{lab}"
            ciname = f"{nm}_ci"
            npaths = int((np.isfinite(A3[f"{ch}_{sname}"]) & (yv == lab)).sum())
            rows.append({"analysis": "E5", "machine": "imi_vm20i", "channel": CH[ch], "setting": SETN[sname], "label": lab, "paths": npaths,
                         "share": round(VA[nm]["value"], 4), "ci_lo": round(VA[ciname]["lo"], 4), "ci_hi": round(VA[ciname]["hi"], 4), "value_name": nm})
e4 = pd.DataFrame(rows)
cols4 = {"analysis": "E1 (clipping) or E5 (off-harmonic dominant line)", "machine, channel, label": "the population",
         "setting": "the clipping criterion (E1) or the indicator setting (E5)", "paths": "paths in the population (E5: spindle speed of at least the comb's lowest speed)",
         "share": "share of those paths meeting the criterion", "ci_lo, ci_hi": "95 percent percentile interval over run resamples drawn within tool families (E5 only)",
         "value_name": "name of the value in the analysis values files (Online Resource 2)"}
with open(f"{OUT}/ESM_4.csv", "w") as fh:
    fh.write(header(4, "shares of clipped paths by machine, sound channel and label (E1) and of paths with an off-harmonic dominant line by channel, indicator setting and label (E5)", cols4))
    e4.to_csv(fh, index=False)
print("ESM_1", d.shape, "ESM_3", e3.shape, "ESM_4", e4.shape)