import numpy as np
import pytest

from delivery.baselines import DistanceLinear, HistoricalAverage
from delivery.data import load_fixture
from delivery.instances import generate_instance
from delivery.nn import TrainConfig, TravelTimeNN, numerical_grad_check, train
from delivery.travel_time import TimeModel


@pytest.fixture(scope="module")
def fixture_data():
    return load_fixture()


def test_backprop_matches_finite_differences():
    assert numerical_grad_check() < 1e-5


def test_fixture_schema_and_filters(fixture_data):
    trips, zones, dist = fixture_data
    for c in ("pu", "do", "hour", "weekday", "duration_min", "distance_mi", "split", "date"):
        assert c in trips.columns
    assert trips.duration_min.between(1, 120).all()
    assert (trips.weekday < 5).all()
    assert len(zones) == 20 and dist.shape == (20, 20)
    # split by day: no date appears in two splits and the test week comes last
    by = trips.groupby("split").date.agg(["min", "max"])
    assert by.loc["train", "max"] < by.loc["val", "min"] <= by.loc["val", "max"] < by.loc["test", "min"]


def test_nn_trains_and_beats_global_mean_on_fixture(fixture_data, tmp_path):
    trips, zones, dist = fixture_data
    tr, va, te = (trips[trips.split == s] for s in ("train", "val", "test"))
    m = train(tr, va, len(zones), dist, TrainConfig(epochs=25, patience=5, batch=128), verbose=False)
    args = [te[c].values for c in ("pu", "do", "hour", "weekday")]
    nn_mae = np.abs(m.predict(*args) - te.duration_min.values).mean()
    mean_mae = np.abs(tr.duration_min.mean() - te.duration_min.values).mean()
    assert nn_mae < mean_mae
    m.save(tmp_path / "m.npz")
    m2 = TravelTimeNN.load(tmp_path / "m.npz")
    assert np.allclose(m.predict(*args), m2.predict(*args))


def test_baselines_fit_and_back_off(fixture_data):
    trips, zones, dist = fixture_data
    tr = trips[trips.split == "train"]
    h = HistoricalAverage(20).fit(tr)
    assert np.isfinite(h.table).all()
    lin = DistanceLinear(dist).fit(tr)
    assert lin.slope_ > 0          # longer typical distance -> longer trips


def test_generator_is_seeded():
    a = generate_instance(20, 7).to_dict()
    b = generate_instance(20, 7).to_dict()
    c = generate_instance(20, 8).to_dict()
    assert a == b and a != c
    assert all(o["latest"] > o["earliest"] for o in a["orders"])


def test_time_model_interpolation_is_fifo_and_exact_at_knots():
    M = np.stack([np.full((2, 2), 10.0 + h) for h in range(24)])
    tm = TimeModel(M, factor=1.2)
    assert tm.tt(0, 1, 8 * 60 + 30) == pytest.approx(18 * 1.2)
    assert tm.tt(0, 1, 9 * 60) == pytest.approx(18.5 * 1.2)
    assert tm.is_fifo()
    assert (tm.lower_bound_matrix(600) <= tm.M[10]).all()