import itertools

import numpy as np
import pytest

from delivery.instances import Order, generate_route_instance
from delivery.search import RouteProblem, brute_force, make_heuristic, plan_van, simulate, solve_route
from delivery.travel_time import TimeModel
from conftest import line_tm


def known_problem():
    # zones on a line, tt = 5*|i-j| + 1; depot zone 0; service 5 min; start 08:00
    stops = [Order(0, 1, 480, 600, 1, 5), Order(1, 2, 480, 600, 1, 5), Order(2, 3, 480, 500, 1, 5)]
    return RouteProblem(stops, depot=0, start=480, tm=line_tm())


# ----------------------------------------------------------------------------- normal cases
@pytest.mark.parametrize("algo", ["astar", "ucs"])
def test_known_answer_small_example(algo):
    # hand-computed: only C(3) first keeps C's window; best is C,B,A = 16+5+6+5+6+5+6 = 49 min
    r = solve_route(known_problem(), algo=algo)
    assert r.sequence == [2, 1, 0]
    assert r.cost == pytest.approx(49.0)
    assert r.feasible and r.stats["expanded"] > 0 and r.stats["generated"] > 0


@pytest.mark.parametrize("seed", range(5))
def test_astar_and_ucs_match_brute_force_on_6_stops(tm, seed):
    inst = generate_route_instance(6, seed, tm, n_zones=tm.Z)
    prob = RouteProblem(inst.orders, inst.depot, inst.vans[0].shift_start, tm)
    bf = brute_force(prob)
    a = solve_route(prob, "astar")
    u = solve_route(prob, "ucs")
    assert bf is not None
    assert a.cost == pytest.approx(bf.cost, abs=1e-6)
    assert u.cost == pytest.approx(bf.cost, abs=1e-6)


def _cost_to_go(prob, cur_zone, clock, remaining):
    best = np.inf
    for perm in itertools.permutations(remaining):
        t, z, ok = clock, cur_zone, True
        for k in perm:
            o = prob.stops[k]
            b = max(t + prob.tm.tt(z, o.zone, t), o.earliest)
            if b > o.latest + 1e-9:
                ok = False
                break
            t, z = b + o.service, o.zone
        if ok:
            t += prob.tm.tt(z, prob.depot, t)
            best = min(best, t - clock)
    return best


def test_admissible_heuristic_never_overestimates(tm):
    rng = np.random.default_rng(1)
    for seed in range(6):
        inst = generate_route_instance(6, seed, tm, n_zones=tm.Z)
        prob = RouteProblem(inst.orders, inst.depot, inst.vans[0].shift_start, tm)
        h = make_heuristic(prob, "admissible")
        for _ in range(15):
            k = int(rng.integers(0, 4))
            done = list(rng.permutation(6)[:k])
            sim = simulate(prob, done)
            cur = done[-1] if done else -1
            clock = sim.departures[-1] if done else prob.start
            mask = sum(1 << d for d in done)
            cz = prob.stops[cur].zone if done else prob.depot
            true = _cost_to_go(prob, cz, clock, [j for j in range(6) if j not in done])
            assert h(cur, mask, clock) <= true + 1e-6


def test_naive_current_hour_heuristic_can_overestimate():
    # traffic eases: 40 min for hours 0-9, 10 min from hour 10. Leaving at 09:40 the
    # interpolated trip takes 40*(5/6) + 10*(1/6) = 35 min, but the naive heuristic reads the
    # hour-9 value (40) -> overestimates -> not admissible. Our lower bound uses the minimum
    # over all later departure times and stays below the true cost.
    M = np.full((24, 2, 2), 10.0)
    M[:10] = 40.0
    clock = 9 * 60 + 40
    prob = RouteProblem([Order(0, 1, 480, 1000, 1, 0)], depot=0, start=clock,
                        tm=TimeModel(M, factor=1.0), return_to_depot=False)
    true = solve_route(prob, "ucs").cost
    assert true == pytest.approx(35.0)
    naive = make_heuristic(prob, "naive")(-1, 0, clock)
    assert naive == pytest.approx(40.0) and naive > true
    assert make_heuristic(prob, "admissible")(-1, 0, clock) <= true


def test_astar_expands_fewer_nodes_than_ucs(tm):
    tot_a = tot_u = 0
    for seed in range(4):
        inst = generate_route_instance(8, seed, tm, n_zones=tm.Z)
        prob = RouteProblem(inst.orders, inst.depot, inst.vans[0].shift_start, tm)
        tot_a += solve_route(prob, "astar").stats["expanded"]
        tot_u += solve_route(prob, "ucs").stats["expanded"]
    assert tot_a < tot_u


def test_greedy_feasible_and_not_better_than_optimal(tm):
    for seed in range(4):
        inst = generate_route_instance(7, seed, tm, n_zones=tm.Z)
        prob = RouteProblem(inst.orders, inst.depot, inst.vans[0].shift_start, tm)
        g = solve_route(prob, "greedy")
        a = solve_route(prob, "astar")
        assert g.feasible
        assert g.cost >= a.cost - 1e-9


# ----------------------------------------------------------------------------- edge cases
def test_edge_empty_van():
    r = plan_van(RouteProblem([], 0, 480, line_tm()))
    assert r.sequence == [] and r.cost == 0 and r.feasible


def test_edge_single_stop_in_depot_zone():
    # intra-zone trip: tt(0,0) = 1 -> arrive 481, serve 5, return 1 -> cost 7
    r = solve_route(RouteProblem([Order(0, 0, 480, 490, 1, 5)], 0, 480, line_tm()))
    assert r.sequence == [0] and r.cost == pytest.approx(7.0)


def test_edge_zero_width_window_hit_exactly_and_boundary_arrival():
    # arrive at zone 2 at 480 + 11 = 491 exactly when the window [491, 491] is open
    r = solve_route(RouteProblem([Order(0, 2, 491, 491, 1, 0)], 0, 480, line_tm(), return_to_depot=False))
    assert r is not None and r.feasible and r.lateness == [0.0]


def test_edge_early_arrival_waits():
    r = solve_route(RouteProblem([Order(0, 1, 600, 610, 1, 5)], 0, 480, line_tm(), return_to_depot=False))
    assert r.waits[0] == pytest.approx(600 - 486)
    assert r.begins[0] == 600


# ----------------------------------------------------------------------------- impossible
def test_impossible_route_returns_none_then_least_late_route():
    # two stops at opposite ends that both close at 08:20: cannot serve both
    stops = [Order(0, 4, 480, 500, 1, 5), Order(1, 0, 480, 500, 1, 5)]
    prob = RouteProblem(stops, 2, 480, line_tm())
    assert solve_route(prob, "astar") is None
    assert brute_force(prob) is None
    soft = plan_van(prob)
    assert not soft.feasible and soft.first_late() is not None
    assert soft.stats.get("soft_fallback")