"""Build data/pedrisk_data.js for the PedRisk website (run once; needs pedrisk_ml.py and the run data)."""
import sys, json, glob, re, os, shutil
import numpy as np, pandas as pd

SCR = "/tmp/claude-0/-home-claude/10e0fb34-63a7-5fb7-a1a4-6d5933d610ba/scratchpad"
DATA = SCR + "/data/pedrisk data"
LABELS = "/mnt/user-data/uploads/conflict_no_conflict_labels_all_validation_runs.csv"
OUT = SCR + "/site"
sys.argv = ["x", DATA, LABELS]
sys.path.insert(0, "/mnt/user-data/outputs")
import pedrisk_ml as pm
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix

runs, X, y, frames, tclose = pm.build()
runs = np.array(runs)
n = len(runs)

# ---- out-of-fold, per-frame scores (repeat 0 of the 8/1/1 rotation)
skf = StratifiedKFold(10, shuffle=True, random_state=0)
folds = [te for _, te in skf.split(X, y)]
oof = {}
runp = {"lr": np.zeros(n), "mlp": np.zeros(n)}
for i in range(10):
    te, va = folds[i], folds[(i + 1) % 10]
    tr = np.concatenate([folds[j] for j in range(10) if j not in (i, (i + 1) % 10)])
    f_lr, C = pm.fit_lr(X[tr], y[tr], X[va], y[va])
    f_mlp, cfg = pm.fit_mlp(X[tr], y[tr], X[va], y[va])
    for k in te:
        t, F = frames[runs[k]]
        oof[runs[k]] = (f_lr(F), f_mlp(F))
        runp["lr"][k] = f_lr(X[k:k + 1])[0]; runp["mlp"][k] = f_mlp(X[k:k + 1])[0]
    print("fold", i, "done", flush=True)


def lead(model):
    rows = []
    for k, r in enumerate(runs):
        t, F = frames[r]
        p = oof[r][0 if model == "lr" else 1]
        hit = np.where(p >= 0.5)[0]
        rows.append((int(y[k]), tclose[r], t[hit[0]] if len(hit) else None))
    c = [(tc, ta) for lab, tc, ta in rows if lab == 1]
    nn = [ta for lab, tc, ta in rows if lab == 0]
    leads = [tc - ta for tc, ta in c if ta is not None]
    return dict(alerted=int(sum(ta is not None for _, ta in c)), n_conflict=len(c),
                before=int(sum(l > 0 for l in leads)),
                median=float(np.median(leads)) if leads else None,
                false_alert_runs=int(sum(a is not None for a in nn)), n_no=len(nn))


results = {
    "cv": {  # 5 repeats of 10-fold (8/1/1), mean +/- sd, from pedrisk_ml.py
        "MLP": dict(precision=[0.56, 0.05], recall=[0.55, 0.05], f1=[0.55, 0.04], accuracy=[0.74, 0.03], auc=[0.77, 0.03]),
        "Logistic regression": dict(precision=[0.59, 0.06], recall=[0.68, 0.09], f1=[0.63, 0.06], accuracy=[0.77, 0.04], auc=[0.80, 0.02]),
        "Distance rule": dict(precision=[0.38, 0.01], recall=[0.92, 0.00], f1=[0.54, 0.01], accuracy=[0.54, 0.01], auc=[0.65, 0.01]),
    },
    "confusion": {"lr": confusion_matrix(y, (runp["lr"] >= .5).astype(int)).tolist(),
                  "mlp": confusion_matrix(y, (runp["mlp"] >= .5).astype(int)).tolist()},
    "lead": {"lr": lead("lr"), "mlp": lead("mlp")},
}
print(json.dumps(results["confusion"]), json.dumps(results["lead"]))

# ---- final models trained on all 41 runs (for the what-if panel)
sc = StandardScaler().fit(X)
lr = LogisticRegression(C=0.1, class_weight="balanced", max_iter=2000).fit(sc.transform(X), y)
mlp = json.load(open("/mnt/user-data/outputs/pedrisk_mlp.json"))
models = {"features": pm.FEATS, "mean": sc.mean_.tolist(), "scale": sc.scale_.tolist(),
          "lr": {"coef": lr.coef_[0].tolist(), "intercept": float(lr.intercept_[0])},
          "mlp": {"W": mlp["W"], "b": mlp["b"], "mean": mlp["mean"], "scale": mlp["scale"]}}

# ---- tracks + per-frame scores
def nice(o):
    return o.replace("VRU_", "").replace("Passenger_Vehicle", "Vehicle").replace("_", " ")

def clean(a):
    return [None if not np.isfinite(v) else round(float(v), 2) for v in a]

R = []
type_counts = {}
for k, r in enumerate(runs):
    t, xy = pm.load_run(f"{DATA}/{r}_GT.csv")
    objs = []
    for o, P in xy.items():
        objs.append({"name": nice(o), "veh": o == "Passenger_Vehicle", "x": clean(P[:, 0]), "y": clean(P[:, 1])})
        if o != "Passenger_Vehicle":
            type_counts.setdefault(nice(o), [0, 0])[int(y[k])] += 1
    p_lr, p_mlp = oof[r]
    R.append({"id": r, "label": int(y[k]), "n": len(t), "t_closest": round(float(tclose[r]), 1),
              "feat": [round(float(v), 3) for v in X[k]], "objs": objs,
              "p_lr": [round(float(v), 3) for v in p_lr], "p_mlp": [round(float(v), 3) for v in p_mlp],
              "run_lr": round(float(runp["lr"][k]), 3), "run_mlp": round(float(runp["mlp"][k]), 3)})
R.sort(key=lambda r: int(r["id"].split("_")[1]))

fr = {"min": X.min(0).round(2).tolist(), "max": X.max(0).round(2).tolist()}
payload = {"runs": R, "models": models, "results": results, "featRange": fr,
           "types": type_counts, "counts": [int((y == 0).sum()), int((y == 1).sum())]}

os.makedirs(OUT + "/data/raw", exist_ok=True)
open(OUT + "/data/pedrisk_data.js", "w").write("window.PEDRISK=" + json.dumps(payload, separators=(",", ":")) + ";")
for f in glob.glob(DATA + "/Run_*_GT.csv"):
    shutil.copy(f, OUT + "/data/raw/")
shutil.copy(LABELS, OUT + "/data/raw/labels.csv")
shutil.copy("/mnt/user-data/outputs/run_features.csv", OUT + "/data/run_features.csv")
shutil.copy("/mnt/user-data/outputs/lead_time_by_run.csv", OUT + "/data/lead_time_by_run.csv")
print("size KB", os.path.getsize(OUT + "/data/pedrisk_data.js") // 1024)