"""Evaluation: prediction metrics and the replay simulation on held-out real trips.

Replay: a plan (fixed stop sequences per van, made with some time model) is driven again,
but each leg i -> j departing at time t takes a duration drawn from REAL test-week trips with
the same (pickup zone, drop-off zone, hour), times the van factor. Pools with fewer than
`min_pool` trips back off to hour +-1, then to all hours of that zone pair. The van still
waits if early. We count late deliveries (service start after the window closes).
Orders a plan could not assign count as late (unserved).
"""
from __future__ import annotations

from collections import defaultdict

import numpy as np

from . import config as C
from .search import simulate
from .integration import van_problem


def metrics(y_true, y_pred) -> dict:
    e = np.asarray(y_pred) - np.asarray(y_true)
    return {"MAE": float(np.abs(e).mean()), "RMSE": float(np.sqrt((e ** 2).mean())),
            "bias": float(e.mean()), "n": int(len(e))}


def metrics_by_band(trips, pred) -> dict:
    bands = {"night 0-5": (0, 6), "morning 6-9": (6, 10), "midday 10-15": (10, 16),
             "evening 16-19": (16, 20), "late 20-23": (20, 24)}
    out = {}
    for name, (a, b) in bands.items():
        m = (trips.hour.values >= a) & (trips.hour.values < b)
        out[name] = metrics(trips.duration_min.values[m], np.asarray(pred)[m])
    return out


class ReplaySampler:
    def __init__(self, test_trips, factor: float = C.VAN_FACTOR, min_pool: int = 3):
        self.factor, self.min_pool = factor, min_pool
        self.cell = defaultdict(list)
        self.pair = defaultdict(list)
        for i, j, h, d in zip(test_trips.pu.values, test_trips["do"].values,
                              test_trips.hour.values, test_trips.duration_min.values):
            self.cell[(int(i), int(j), int(h))].append(float(d))
            self.pair[(int(i), int(j))].append(float(d))
        self.cell = {k: np.array(v) for k, v in self.cell.items()}
        self.pair = {k: np.array(v) for k, v in self.pair.items()}
        self.backoffs = defaultdict(int)

    def pool(self, i, j, t):
        h = int(t // 60) % 24
        p = self.cell.get((i, j, h))
        if p is not None and len(p) >= self.min_pool:
            return p
        near = [self.cell.get((i, j, hh)) for hh in (h - 1, h, h + 1)]
        near = [x for x in near if x is not None]
        if near and sum(len(x) for x in near) >= self.min_pool:
            self.backoffs["hour+-1"] += 1
            return np.concatenate(near)
        if (i, j) in self.pair:
            self.backoffs["all-hours"] += 1
            return self.pair[(i, j)]
        self.backoffs["reverse-pair"] += 1
        return self.pair.get((j, i), np.array([10.0]))

    def travel_fn(self, rng):
        def f(i, j, t):
            p = self.pool(int(i), int(j), t)
            return float(p[rng.integers(len(p))]) * self.factor
        return f


def replay(inst, result, sampler: ReplaySampler, tm, n_reps: int = 30, seed: int = 0) -> dict:
    """Late deliveries when the plan's sequences are driven with real held-out durations."""
    rng = np.random.default_rng(seed)
    served = {i for r in result.routes.values() for i in r.sequence}
    unserved = len(inst.orders) - len(served)
    lates, minutes = [], []
    for _ in range(n_reps):
        fn = sampler.travel_fn(rng)
        n_late, tot = unserved, 0.0
        for v in inst.vans:
            r = result.routes.get(v.id)
            if r is None or not r.sequence:
                continue
            prob = van_problem(inst, tm, v, r.sequence)
            sim = simulate(prob, list(range(len(r.sequence))), travel=fn)
            n_late += sim.n_late
            tot += sum(sim.lateness)
        lates.append(n_late); minutes.append(tot)
    return {"late_mean": float(np.mean(lates)), "late_std": float(np.std(lates)),
            "late_minutes_mean": float(np.mean(minutes)), "unserved": unserved,
            "late_frac": float(np.mean(lates)) / max(1, len(inst.orders))}