"""Hand-written numpy MLP that predicts taxi trip duration (minutes).

Inputs (all available at planning time, no leakage of the actual trip):
  pickup zone   one-hot (Z)  -> the first layer's weights on it are a learned zone embedding
  drop-off zone one-hot (Z)  -> second embedding
  hour          one-hot (24) + sin/cos of the hour (smoothness hint)
  weekday       one-hot (5)
  zone-pair mean distance (train split only), standardized, plus its log

Network: D -> 64 -> 32 -> 1, ReLU, Huber loss on the standardized target,
Adam, mini-batches, L2 weight decay, early stopping on validation MAE.
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np


# --------------------------------------------------------------------------- features
class Featurizer:
    def __init__(self, n_zones: int, pair_dist: np.ndarray):
        self.Z = n_zones
        self.pair_dist = np.asarray(pair_dist, dtype=float)
        d = self.pair_dist.ravel()
        self.d_mu, self.d_sd = float(d.mean()), float(d.std() + 1e-9)
        ld = np.log1p(d)
        self.ld_mu, self.ld_sd = float(ld.mean()), float(ld.std() + 1e-9)

    @property
    def dim(self) -> int:
        return 2 * self.Z + 24 + 2 + 5 + 2

    def transform(self, pu, do, hour, weekday) -> np.ndarray:
        pu, do = np.asarray(pu, int), np.asarray(do, int)
        hour, weekday = np.asarray(hour, int), np.asarray(weekday, int)
        n, Z = len(pu), self.Z
        X = np.zeros((n, self.dim), dtype=np.float32)
        r = np.arange(n)
        X[r, pu] = 1.0
        X[r, Z + do] = 1.0
        X[r, 2 * Z + hour] = 1.0
        o = 2 * Z + 24
        X[:, o] = np.sin(2 * np.pi * hour / 24)
        X[:, o + 1] = np.cos(2 * np.pi * hour / 24)
        X[r, o + 2 + weekday] = 1.0
        d = self.pair_dist[pu, do]
        X[:, o + 7] = (d - self.d_mu) / self.d_sd
        X[:, o + 8] = (np.log1p(d) - self.ld_mu) / self.ld_sd
        return X

    def state(self) -> dict:
        return {"Z": self.Z, "pair_dist": self.pair_dist}

    @classmethod
    def from_state(cls, s) -> "Featurizer":
        return cls(int(s["Z"]), s["pair_dist"])


# --------------------------------------------------------------------------- network
def huber(r: np.ndarray, delta: float):
    a = np.abs(r)
    loss = np.where(a <= delta, 0.5 * r * r, delta * (a - 0.5 * delta))
    grad = np.where(a <= delta, r, delta * np.sign(r))
    return loss.mean(), grad / len(r)


class MLP:
    def __init__(self, sizes, seed: int = 0):
        rng = np.random.default_rng(seed)
        self.sizes = list(sizes)
        self.W, self.b = [], []
        for a, b in zip(sizes[:-1], sizes[1:]):
            self.W.append(rng.normal(0, np.sqrt(2.0 / a), size=(a, b)).astype(np.float64))
            self.b.append(np.zeros(b))

    def params(self):
        return self.W + self.b

    def forward(self, X):
        acts = [X]
        h = X
        for k, (W, b) in enumerate(zip(self.W, self.b)):
            z = h @ W + b
            h = np.maximum(z, 0.0) if k < len(self.W) - 1 else z
            acts.append(h)
        return h[:, 0], acts

    def backward(self, acts, dout):
        """dout: dL/d(output) shape (n,). Returns grads aligned with params()."""
        gW = [None] * len(self.W)
        gb = [None] * len(self.b)
        g = dout[:, None]
        for k in reversed(range(len(self.W))):
            gW[k] = acts[k].T @ g
            gb[k] = g.sum(axis=0)
            if k > 0:
                g = (g @ self.W[k].T) * (acts[k] > 0)
        return gW + gb


class Adam:
    def __init__(self, params, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, weight_decay=0.0):
        self.p = params
        self.lr, self.b1, self.b2, self.eps, self.wd = lr, b1, b2, eps, weight_decay
        self.m = [np.zeros_like(x) for x in params]
        self.v = [np.zeros_like(x) for x in params]
        self.t = 0

    def step(self, grads, n_weight_mats: int):
        self.t += 1
        for i, (p, g) in enumerate(zip(self.p, grads)):
            if self.wd and i < n_weight_mats:          # decay weights, not biases
                g = g + self.wd * p
            self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
            self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * g * g
            mh = self.m[i] / (1 - self.b1 ** self.t)
            vh = self.v[i] / (1 - self.b2 ** self.t)
            p -= self.lr * mh / (np.sqrt(vh) + self.eps)


# --------------------------------------------------------------------------- model
@dataclass
class TrainConfig:
    hidden: tuple = (64, 32)
    lr: float = 1e-3
    batch: int = 512
    epochs: int = 40
    patience: int = 5
    weight_decay: float = 1e-5
    huber_delta: float = 1.0
    seed: int = 42


@dataclass
class TravelTimeNN:
    feat: Featurizer
    net: MLP
    y_mu: float
    y_sd: float
    history: list = field(default_factory=list)

    def predict(self, pu, do, hour, weekday) -> np.ndarray:
        X = self.feat.transform(pu, do, hour, weekday)
        out, _ = self.net.forward(X)
        return np.clip(out * self.y_sd + self.y_mu, 1.0, 120.0)

    # ---- persistence
    def save(self, path: Path):
        arrs = {f"W{i}": w for i, w in enumerate(self.net.W)}
        arrs.update({f"b{i}": b for i, b in enumerate(self.net.b)})
        np.savez_compressed(path, sizes=np.array(self.net.sizes), y_mu=self.y_mu, y_sd=self.y_sd,
                            Z=self.feat.Z, pair_dist=self.feat.pair_dist,
                            history=np.array(self.history, dtype=float), **arrs)

    @classmethod
    def load(cls, path: Path) -> "TravelTimeNN":
        f = np.load(path)
        sizes = [int(s) for s in f["sizes"]]
        net = MLP(sizes)
        net.W = [f[f"W{i}"] for i in range(len(sizes) - 1)]
        net.b = [f[f"b{i}"] for i in range(len(sizes) - 1)]
        feat = Featurizer(int(f["Z"]), f["pair_dist"])
        return cls(feat, net, float(f["y_mu"]), float(f["y_sd"]), f["history"].tolist())


def train(trips_train, trips_val, n_zones: int, pair_dist: np.ndarray,
          cfg: TrainConfig = TrainConfig(), verbose: bool = True) -> TravelTimeNN:
    feat = Featurizer(n_zones, pair_dist)
    cols = ["pu", "do", "hour", "weekday"]
    Xtr = feat.transform(*[trips_train[c].values for c in cols]).astype(np.float64)
    Xva = feat.transform(*[trips_val[c].values for c in cols]).astype(np.float64)
    ytr = trips_train.duration_min.values.astype(float)
    yva = trips_val.duration_min.values.astype(float)
    y_mu, y_sd = float(ytr.mean()), float(ytr.std())
    ttr = (ytr - y_mu) / y_sd

    net = MLP([feat.dim, *cfg.hidden, 1], seed=cfg.seed)
    opt = Adam(net.params(), lr=cfg.lr, weight_decay=cfg.weight_decay)
    rng = np.random.default_rng(cfg.seed)
    model = TravelTimeNN(feat, net, y_mu, y_sd)
    best, best_state, bad = np.inf, None, 0
    n = len(Xtr)
    for ep in range(1, cfg.epochs + 1):
        t0 = time.time()
        perm = rng.permutation(n)
        tot = 0.0
        for s in range(0, n, cfg.batch):
            idx = perm[s:s + cfg.batch]
            out, acts = net.forward(Xtr[idx])
            loss, dout = huber(out - ttr[idx], cfg.huber_delta)
            tot += loss * len(idx)
            opt.step(net.backward(acts, dout), len(net.W))
        pva = np.clip(net.forward(Xva)[0] * y_sd + y_mu, 1, 120)
        va_mae = float(np.abs(pva - yva).mean())
        model.history.append((ep, tot / n, va_mae))
        if verbose:
            print(f"  epoch {ep:2d}  train huber {tot/n:.4f}  val MAE {va_mae:.3f} min  ({time.time()-t0:.1f}s)")
        if va_mae < best - 1e-4:
            best, bad = va_mae, 0
            best_state = ([w.copy() for w in net.W], [b.copy() for b in net.b])
        else:
            bad += 1
            if bad >= cfg.patience:
                if verbose:
                    print(f"  early stop at epoch {ep} (best val MAE {best:.3f})")
                break
    net.W, net.b = best_state
    # Adam holds references to the old arrays; the model itself now uses the best weights
    return model


def numerical_grad_check(seed: int = 0, eps: float = 1e-6) -> float:
    """Max relative error between backprop and finite-difference gradients (for tests)."""
    rng = np.random.default_rng(seed)
    net = MLP([7, 5, 4, 1], seed=seed)
    X = rng.normal(size=(9, 7))
    y = rng.normal(size=9)

    def L():
        return huber(net.forward(X)[0] - y, 1.0)[0]

    out, acts = net.forward(X)
    _, dout = huber(out - y, 1.0)
    grads = net.backward(acts, dout)
    worst = 0.0
    for p, g in zip(net.params(), grads):
        it = np.nditer(p, flags=["multi_index"])
        for _ in it:
            i = it.multi_index
            old = p[i]
            p[i] = old + eps; lp = L()
            p[i] = old - eps; lm = L()
            p[i] = old
            num = (lp - lm) / (2 * eps)
            worst = max(worst, abs(num - g[i]) / max(1e-8, abs(num) + abs(g[i])))
    return worst