"""Zone-to-zone predicted travel-time matrices by hour, and how the planner reads them.

A matrix M has shape (24, Z, Z): M[h, i, j] = predicted taxi minutes from zone i to zone j
for a departure during hour h (one weekday). The van model multiplies by VAN_FACTOR (1.2).

Departure-time interpolation: the hourly value is placed at the middle of the hour
(h*60+30) and interpolated linearly in between. Linear interpolation keeps travel time
continuous and makes the network FIFO ("leaving later never makes you arrive earlier")
as long as |M[h+1]-M[h]| < 60 min, which `is_fifo` checks. FIFO + waiting allowed is what
makes "earliest clock time dominates" valid in the search (see search.py).
"""
from __future__ import annotations

import numpy as np

from . import config as C


def build_matrix(predictor, n_zones: int, weekday: int) -> np.ndarray:
    Z = n_zones
    h, i, j = np.meshgrid(np.arange(24), np.arange(Z), np.arange(Z), indexing="ij")
    pred = predictor.predict(i.ravel(), j.ravel(), h.ravel(), np.full(h.size, weekday))
    return np.asarray(pred, float).reshape(24, Z, Z)


def build_all_weekdays(predictor, n_zones: int) -> np.ndarray:
    return np.stack([build_matrix(predictor, n_zones, w) for w in range(5)])


class TimeModel:
    """Travel-time oracle used by the CSP, the search and the simulator."""

    def __init__(self, matrix: np.ndarray, factor: float = C.VAN_FACTOR, name: str = "model"):
        M = np.asarray(matrix, float)
        if M.ndim == 2:                      # constant matrix -> same for every hour
            M = np.repeat(M[None], 24, axis=0)
        self.M = M * factor
        self.factor, self.name = factor, name
        self.Z = M.shape[1]
        # suffix minimum over hour knots: lower bound on travel time for any departure >= t
        self.suffix_min = np.minimum.accumulate(self.M[::-1], axis=0)[::-1]
        self.global_min = self.M.min(axis=0)

    @staticmethod
    def _knot(t: float):
        x = (t - 30.0) / 60.0
        if x <= 0:
            return 0, 0.0
        if x >= 23:
            return 23, 0.0
        h0 = int(np.floor(x))
        return h0, x - h0

    def tt(self, i: int, j: int, t: float) -> float:
        h0, w = self._knot(t)
        if w == 0.0:
            return float(self.M[h0, i, j])
        return float((1 - w) * self.M[h0, i, j] + w * self.M[h0 + 1, i, j])

    def lower_bound_matrix(self, t: float) -> np.ndarray:
        """min over all departure times >= t of tt(., ., .) -- admissible edge costs."""
        h0, _ = self._knot(t)
        return self.suffix_min[h0]

    def is_fifo(self) -> bool:
        return bool(np.all(np.diff(self.M, axis=0) > -60.0))