"""Train the numpy MLP, fit the two baselines, evaluate on the held-out test week,
and export the predicted zone-to-zone time matrices by hour.

    python train_nn.py
"""
import json

import numpy as np

from delivery import config as C
from delivery import plots as P
from delivery.baselines import DistanceLinear, HistoricalAverage, NNPredictor
from delivery.data import load_processed
from delivery.eval import metrics, metrics_by_band
from delivery.nn import TrainConfig, train
from delivery.travel_time import TimeModel, build_all_weekdays


def main():
    trips, zones, dist = load_processed()
    Z = len(zones)
    tr, va, te = (trips[trips.split == s] for s in ("train", "val", "test"))
    print(f"train {len(tr):,}  val {len(va):,}  test {len(te):,}  zones {Z}")

    print("training NN (numpy MLP, Adam, early stopping)")
    cfg = TrainConfig(seed=C.SEED)
    model = train(tr, va, Z, dist, cfg)
    C.MODELS.mkdir(exist_ok=True)
    model.save(C.MODELS / "mlp.npz")

    hist = HistoricalAverage(Z).fit(tr)
    lin = DistanceLinear(dist).fit(tr)
    nn = NNPredictor(model)
    args = [te[c].values for c in ("pu", "do", "hour", "weekday")]
    preds = {"NN (MLP)": nn.predict(*args), "hist-avg (pair,hour)": hist.predict(*args),
             "distance-linear": lin.predict(*args)}
    preds_all = dict(preds, **{"global mean": np.full(len(te), tr.duration_min.mean())})
    y = te.duration_min.values
    res = {"test": {k: metrics(y, p) for k, p in preds_all.items()},
           "by_band": {k: metrics_by_band(te, p) for k, p in preds.items()},
           "train_config": cfg.__dict__, "epochs_run": len(model.history),
           "best_val_mae": float(min(h[2] for h in model.history)),
           "distance_linear_coef": {"intercept": float(lin.intercept_), "slope": float(lin.slope_)}}
    # rare night cells: where hist-avg had to back off
    rare = hist.counts[te.pu.values, te["do"].values, te.hour.values] < hist.min_count
    res["rare_cells"] = {"test_trips_in_rare_cells": int(rare.sum()),
                         **{k: metrics(y[rare], p[rare]) for k, p in preds.items()}} if rare.any() else {}
    print("\n test-week results (minutes)")
    print(f"  {'model':<24}{'MAE':>8}{'RMSE':>8}{'bias':>8}")
    for k, m in res["test"].items():
        print(f"  {k:<24}{m['MAE']:8.3f}{m['RMSE']:8.3f}{m['bias']:8.3f}")

    C.OUTPUTS.mkdir(exist_ok=True)
    (C.OUTPUTS / "nn_results.json").write_text(json.dumps(res, indent=2))
    (C.ROOT / "report" / "results").mkdir(parents=True, exist_ok=True)
    (C.ROOT / "report" / "results" / "nn_results.json").write_text(json.dumps(res, indent=2))

    # time matrices (taxi minutes; the planner multiplies by VAN_FACTOR)
    mats = {name: build_all_weekdays(p, Z) for name, p in
            (("nn", nn), ("hist", hist), ("dist", lin))}
    np.savez_compressed(C.MODELS / "time_matrices.npz", **mats)
    print("FIFO check (|M[h+1]-M[h]| < 60):", {k: TimeModel(m[1]).is_fifo() for k, m in mats.items()})
    with open(C.OUTPUTS / "time_matrix_nn_tuesday.csv", "w") as f:
        f.write("hour,pu_idx,do_idx,pu_zone,do_zone,minutes\n")
        names = zones.sort_values("idx").zone.tolist()
        for h in range(24):
            for i in range(Z):
                for j in range(Z):
                    f.write(f"{h},{i},{j},\"{names[i]}\",\"{names[j]}\",{mats['nn'][1, h, i, j]:.2f}\n")

    P.learning_curve(model.history, C.OUTPUTS / "nn_learning_curve.png")
    P.mae_by_hour(te, preds, C.OUTPUTS / "nn_mae_by_hour.png")
    P.pred_vs_actual(y, preds["NN (MLP)"], C.OUTPUTS / "nn_pred_vs_actual.png")
    P.matrix_heatmaps(mats["nn"][1], zones, C.OUTPUTS / "time_matrix_nn.png")
    print("saved models/mlp.npz, models/time_matrices.npz and figures in outputs/")


if __name__ == "__main__":
    main()