"""
PedRisk: predict conflict / no conflict for each FHWA validation run from ground-truth tracks.
Models: MLP (main), logistic regression (slide), distance-threshold rule (baseline).
Evaluation: run-level, 8/1/1 fold rotation (80% train / 10% validation / 10% test), repeated.
Usage: python pedrisk_ml.py <folder with Run_XX_GT.csv> <labels.csv>
"""
import sys, glob, re, json, warnings
import numpy as np, pandas as pd
from sklearn.neural_network import MLPClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import (precision_score, recall_score, f1_score, accuracy_score,
                             roc_auc_score, log_loss, confusion_matrix)
warnings.filterwarnings("ignore")

DATA = sys.argv[1] if len(sys.argv) > 1 else "/tmp/pd/pedrisk data"
LABELS = sys.argv[2] if len(sys.argv) > 2 else "/mnt/user-data/uploads/conflict_no_conflict_labels_all_validation_runs.csv"
DT = 0.1
FAR = 50.0  # value used when no vehicle is tracked

FEATS = ["min_dist", "min_dcpa", "veh_max_speed", "veh_speed_at_closest", "max_closing",
         "min_ttc", "max_decel", "n_vru", "has_child", "has_fast_wheeled", "has_assistive"]
FAST = ("Motorized_Bicycle", "Manual_Bicycle", "Manual_Scooter", "Skateboard")
ASSIST = ("Walker", "Cane", "Crutches", "Wheelchair", "Stroller")


def bdiff(a, k=4):
    """Causal velocity: backward difference over k frames (NaN where untracked)."""
    out = np.full_like(a, np.nan)
    out[k:] = (a[k:] - a[:-k]) / (k * DT)
    return out


def load_run(path):
    d = pd.read_csv(path)
    objs = sorted({c.rsplit("_", 1)[0] for c in d.columns if c.endswith("_xctr")})
    xy = {o: d[[o + "_xctr", o + "_yctr"]].values.astype(float) for o in objs}
    return d.Time.values - d.Time.values[0], xy


def cum_features(t, xy):
    """Return (n_frames x n_features) array of features computed ONLY from data up to each frame."""
    n = len(t)
    veh = xy.get("Passenger_Vehicle")
    vrus = {o: v for o, v in xy.items() if o != "Passenger_Vehicle"}
    names = list(vrus)
    static = [len(vrus),
              float(any("Child" in o for o in names)),
              float(any(any(f in o for f in FAST) for o in names)),
              float(any(any(a in o for a in ASSIST) for o in names))]
    F = np.zeros((n, len(FEATS)))
    F[:, 7:] = static
    if veh is None or not np.isfinite(veh[:, 0]).any() or not vrus:
        F[:, 0] = F[:, 1] = FAR; F[:, 5] = 20.0
        return F
    vv = np.c_[bdiff(veh[:, 0]), bdiff(veh[:, 1])]
    speed = np.hypot(vv[:, 0], vv[:, 1])
    speed = np.where(speed > 35, np.nan, speed)           # drop tracking glitches
    acc = bdiff(speed)
    dist_all = np.full(n, np.inf); dcpa_all = np.full(n, np.inf)
    close_all = np.zeros(n); ttc_all = np.full(n, 20.0)
    for o, P in vrus.items():
        pv = np.c_[bdiff(P[:, 0]), bdiff(P[:, 1])]
        r = P - veh; vr = pv - vv
        dist = np.hypot(r[:, 0], r[:, 1])
        rate = np.einsum("ij,ij->i", r, vr) / np.maximum(dist, 1e-6)   # d(dist)/dt
        closing = np.where(np.isfinite(rate), np.maximum(-rate, 0), 0)
        tcpa = -np.einsum("ij,ij->i", r, vr) / np.maximum(np.einsum("ij,ij->i", vr, vr), 1e-6)
        tcpa = np.clip(tcpa, 0, 5)
        dcpa = np.hypot(*(r + vr * tcpa[:, None]).T)
        moving = speed > 0.5
        dcpa = np.where(moving & np.isfinite(dcpa), dcpa, np.inf)
        ttc = np.where((closing > 0.3) & (dist < 40), dist / np.maximum(closing, 1e-6), 20.0)
        dist_all = np.fmin(dist_all, np.where(np.isfinite(dist), dist, np.inf))
        dcpa_all = np.fmin(dcpa_all, dcpa)
        close_all = np.fmax(close_all, np.where(dist < 15, closing, 0))
        ttc_all = np.fmin(ttc_all, np.minimum(ttc, 20.0))
    cm_dist = np.minimum.accumulate(dist_all)
    F[:, 0] = np.where(np.isfinite(cm_dist), cm_dist, FAR)
    cm_dcpa = np.minimum.accumulate(dcpa_all)
    F[:, 1] = np.where(np.isfinite(cm_dcpa), cm_dcpa, FAR)
    F[:, 2] = np.fmax.accumulate(np.nan_to_num(speed, nan=0.0))
    # vehicle speed at the moment of the running minimum distance
    best, sp_at = np.inf, 0.0
    for k in range(n):
        if np.isfinite(dist_all[k]) and dist_all[k] <= best:
            best = dist_all[k]; sp_at = 0.0 if not np.isfinite(speed[k]) else speed[k]
        F[k, 3] = sp_at
    F[:, 4] = np.maximum.accumulate(close_all)
    F[:, 5] = np.minimum.accumulate(ttc_all)
    F[:, 6] = np.clip(np.fmax.accumulate(np.nan_to_num(-acc, nan=0.0)), 0, 10)
    return F


def build():
    lab = pd.read_csv(LABELS); lab.columns = ["run", "label"]
    lab = dict(zip(lab.run, (lab.label == "conflict").astype(int)))
    runs, X, y, frames, tclose = [], [], [], {}, {}
    for f in sorted(glob.glob(f"{DATA}/Run_*_GT.csv")):
        run = re.search(r"(Run_\d+)_GT", f).group(1)
        t, xy = load_run(f)
        F = cum_features(t, xy)
        runs.append(run); X.append(F[-1]); y.append(lab[run]); frames[run] = (t, F)
        # time of overall closest approach (for lead-time analysis)
        tclose[run] = t[int(np.argmin(np.where(np.isfinite(F[:, 0]), F[:, 0], 1e9) + 0))] if False else None
        k = np.where(F[:, 0] <= F[-1, 0] + 1e-9)[0][0]
        tclose[run] = t[k]
    return runs, np.array(X), np.array(y), frames, tclose


# ---------------------------------------------------------------- models
def balance(X, y, seed=0):
    rng = np.random.RandomState(seed)
    pos, neg = np.where(y == 1)[0], np.where(y == 0)[0]
    big, small = (neg, pos) if len(neg) > len(pos) else (pos, neg)
    idx = np.concatenate([big, rng.choice(small, len(big), replace=True)])
    return X[idx], y[idx]


def fit_mlp(Xtr, ytr, Xva, yva, seeds=3):
    sc = StandardScaler().fit(Xtr)
    Xb, yb = balance(sc.transform(Xtr), ytr)
    best = None
    for h in [(4,), (8,), (8, 4)]:
        for al in [1e-2, 1e-1, 1.0]:
            ps = np.mean([MLPClassifier(h, alpha=al, max_iter=2000, random_state=s)
                          .fit(Xb, yb).predict_proba(sc.transform(Xva))[:, 1] for s in range(seeds)], 0)
            ll = log_loss(yva, np.clip(ps, 1e-3, 1 - 1e-3), labels=[0, 1])
            if best is None or ll < best[0]:
                best = (ll, h, al)
    _, h, al = best
    models = [MLPClassifier(h, alpha=al, max_iter=2000, random_state=s).fit(Xb, yb) for s in range(5)]
    return (lambda X: np.mean([m.predict_proba(sc.transform(X))[:, 1] for m in models], 0)), (h, al)


def fit_lr(Xtr, ytr, Xva, yva):
    sc = StandardScaler().fit(Xtr)
    best = None
    for C in [0.01, 0.1, 1, 10]:
        m = LogisticRegression(C=C, class_weight="balanced", max_iter=2000).fit(sc.transform(Xtr), ytr)
        ll = log_loss(yva, np.clip(m.predict_proba(sc.transform(Xva))[:, 1], 1e-3, 1 - 1e-3), labels=[0, 1])
        if best is None or ll < best[0]:
            best = (ll, m, C)
    return (lambda X: best[1].predict_proba(sc.transform(X))[:, 1]), best[2]


def fit_rule(Xtr, ytr):
    """Baseline: conflict if min distance < threshold; threshold maximises F1 on training runs."""
    cands = np.unique(Xtr[:, 0])
    th = max(cands, key=lambda c: f1_score(ytr, (Xtr[:, 0] <= c).astype(int), zero_division=0))
    return (lambda X: (X[:, 0] <= th).astype(float)), th


def metrics(y, p, thr=0.5):
    yh = (p >= thr).astype(int)
    out = dict(precision=precision_score(y, yh, zero_division=0), recall=recall_score(y, yh),
               f1=f1_score(y, yh, zero_division=0), accuracy=accuracy_score(y, yh))
    out["auc"] = roc_auc_score(y, p) if len(np.unique(p)) > 1 else 0.5
    return out


def main():
    runs, X, y, frames, tclose = build()
    runs = np.array(runs)
    print(f"{len(runs)} runs, {y.sum()} conflict / {(1 - y).sum()} no conflict")
    df = pd.DataFrame(X, columns=FEATS); df.insert(0, "run", runs); df["conflict"] = y
    df.to_csv("/mnt/user-data/outputs/run_features.csv", index=False)
    print("\nFeature medians by label:")
    print(df.groupby("conflict")[FEATS].median().round(2).T.to_string())

    REPEATS = 5
    res = {m: [] for m in ["MLP", "LogReg", "Rule"]}
    pooled = {m: np.zeros((REPEATS, len(y))) for m in res}
    chosen, lead_info = [], []
    for rep in range(REPEATS):
        skf = StratifiedKFold(10, shuffle=True, random_state=rep)
        folds = [te for _, te in skf.split(X, y)]
        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_mlp, cfg = fit_mlp(X[tr], y[tr], X[va], y[va])
            f_lr, C = fit_lr(X[tr], y[tr], X[va], y[va])
            f_rl, th = fit_rule(X[tr], y[tr])
            pooled["MLP"][rep, te] = f_mlp(X[te]); pooled["LogReg"][rep, te] = f_lr(X[te])
            pooled["Rule"][rep, te] = f_rl(X[te])
            if rep == 0:
                chosen.append(cfg)
                for k in te:   # causal replay of each test run: alert when running MLP probability >= 0.5
                    t, F = frames[runs[k]]
                    p = f_mlp(F)
                    hit = np.where(p >= 0.5)[0]
                    lead_info.append((runs[k], y[k], tclose[runs[k]], t[hit[0]] if len(hit) else None))
        for m in res:
            res[m].append(metrics(y, pooled[m][rep]))
    print("\nTest-set performance, all 41 runs (each run is a test run once per repeat); mean +/- sd over", REPEATS, "repeats")
    tab = []
    for m in res:
        r = pd.DataFrame(res[m])
        tab.append({"model": m, **{c: f"{r[c].mean():.2f} +/- {r[c].std():.2f}" for c in r}})
    print(pd.DataFrame(tab).to_string(index=False))
    cm = confusion_matrix(y, (pooled["MLP"][0] >= .5).astype(int))
    print("\nMLP confusion matrix, repeat 0 (rows true no-conflict/conflict, cols predicted):\n", cm)
    mean_p = pooled["MLP"].mean(0)
    wrong = [(runs[k], int(y[k]), round(float(mean_p[k]), 2)) for k in range(len(y)) if (mean_p[k] >= .5) != y[k]]
    print("Runs the MLP gets wrong on average (run, true, mean prob):", wrong)
    print("Most chosen MLP configs (hidden, alpha):", pd.Series([str(c) for c in chosen]).value_counts().to_dict())

    L = pd.DataFrame(lead_info, columns=["run", "conflict", "t_closest", "t_alert"])
    L["lead"] = L.t_closest - L.t_alert
    c = L[L.conflict == 1]; n = L[L.conflict == 0]
    print("\nWarning lead time (MLP replayed causally, alert at probability >= 0.5):")
    print(f"  conflict runs alerted: {c.t_alert.notna().sum()}/{len(c)}; "
          f"alert before closest approach: {(c.lead > 0).sum()}/{len(c)}; "
          f"median lead {c.lead.dropna().median():.1f} s")
    print(f"  no-conflict runs that raised an alert at any time: {n.t_alert.notna().sum()}/{len(n)}")
    L.to_csv("/mnt/user-data/outputs/lead_time_by_run.csv", index=False)

    # final model on all runs, exported for the website
    sc = StandardScaler().fit(X)
    Xb, yb = balance(sc.transform(X), y)
    m = MLPClassifier((8,), alpha=0.1, max_iter=3000, random_state=0).fit(Xb, yb)
    json.dump({"features": FEATS, "mean": sc.mean_.tolist(), "scale": sc.scale_.tolist(),
               "W": [w.tolist() for w in m.coefs_], "b": [b.tolist() for b in m.intercepts_]},
              open("/mnt/user-data/outputs/pedrisk_mlp.json", "w"))


if __name__ == "__main__":
    main()