"""Stop ordering for one van: UCS, Greedy best-first and A*, plus brute force.

State       (current stop, set of visited stops, clock time)
Action      drive to an unvisited stop j, wait if early, serve it
Cost        g = clock - shift start = predicted travel + waiting + service
            (service is a constant, so this is "travel + waiting" up to a constant);
            with return_to_depot the final leg back to the depot is included.
Hard mode   arriving after a window closes is not allowed (successor pruned).
Soft mode   late service allowed, g += LATE_PENALTY * minutes late (used only to find the
            least-late route when no on-time route exists, to locate the "first late pair").

Duplicate detection / dominance
  Travel times are FIFO (see travel_time.py) and waiting is allowed, so for the same
  (current stop, visited set) a state with an earlier clock (and, in soft mode, no more
  lateness) can do everything the later one can, at least as cheaply. We keep only the
  Pareto-optimal (clock, lateness) labels per (stop, visited) -- in hard mode a single label.

A* heuristic (admissible, proof in report section 4.2)
  LB[u,k] = min over every departure time >= now of the predicted time u -> k.
  h_sum  = sum over unvisited k of ( min_{u in {cur} U unvisited, u != k} LB[u,k] + service_k )
           + min_{u in unvisited} LB[u, depot]                  (if returning)
  Every unvisited stop is entered by exactly one distinct leg, every leg costs at least its
  LB, waiting is >= 0, so h_sum <= true remaining cost.
  h_wait = max over unvisited k of ( max(e_k, now + min_u LB[u,k]) + service_k ) - now (+ return LB)
  Each stop k is still to be served, no earlier than its window opens: also a lower bound.
  h = max(h_sum, h_wait) is admissible (max of admissible heuristics).
  The "naive" version (cheapest incoming edge at the CURRENT hour) is NOT admissible when
  traffic gets lighter later in the day or the van must wait for a window - we keep it only
  to demonstrate that in the tests.
  A* re-opens states when a cheaper path is found, so admissibility alone gives optimality.
"""
from __future__ import annotations

import heapq
import itertools
import time
from dataclasses import dataclass, field

import numpy as np

LATE_PENALTY = 1000.0
DEPOT = -1


@dataclass
class Route:
    van: int
    sequence: list                 # order ids in visiting order
    arrivals: list
    begins: list                   # service start (after waiting)
    departures: list
    waits: list
    lateness: list
    finish: float                  # back at depot (or end of last service)
    cost: float                    # finish - start (+ penalty in soft mode)
    feasible: bool
    stats: dict = field(default_factory=dict)

    @property
    def n_late(self):
        return int(sum(1 for x in self.lateness if x > 1e-9))

    def first_late(self):
        for k, x in enumerate(self.lateness):
            if x > 1e-9:
                return k
        return None


@dataclass
class RouteProblem:
    stops: list          # Order objects
    depot: int
    start: float
    tm: object           # TimeModel
    return_to_depot: bool = True
    shift_end: float = float("inf")
    van: int = 0


# --------------------------------------------------------------------------- simulation
def simulate(problem: RouteProblem, seq_local, travel=None) -> Route:
    """Drive a fixed sequence (local indices). `travel(i_zone, j_zone, t)` defaults to the model."""
    travel = travel or problem.tm.tt
    t, cur = problem.start, problem.depot
    arr, beg, dep, wai, late = [], [], [], [], []
    for k in seq_local:
        o = problem.stops[k]
        a = t + travel(cur, o.zone, t)
        b = max(a, o.earliest)
        arr.append(a); beg.append(b); wai.append(b - a); late.append(max(0.0, b - o.latest))
        t = b + o.service
        dep.append(t)
        cur = o.zone
    if problem.return_to_depot and seq_local:
        t = t + travel(cur, problem.depot, t)
    over = max(0.0, t - problem.shift_end)
    feasible = all(x <= 1e-9 for x in late) and over <= 1e-9
    cost = t - problem.start + LATE_PENALTY * (sum(late) + over)
    return Route(problem.van, [problem.stops[k].id for k in seq_local], arr, beg, dep, wai, late,
                 t, cost, feasible)


def brute_force(problem: RouteProblem, hard: bool = True) -> Route | None:
    n = len(problem.stops)
    best, t0, count = None, time.perf_counter(), 0
    for perm in itertools.permutations(range(n)):
        count += 1
        r = simulate(problem, perm)
        if hard and not r.feasible:
            continue
        if best is None or r.cost < best.cost - 1e-9:
            best = r
    if best is not None:
        best.stats = {"algo": "brute-force", "permutations": count,
                      "time_s": time.perf_counter() - t0}
    return best


# --------------------------------------------------------------------------- heuristics
def make_heuristic(problem: RouteProblem, kind: str):
    stops, tm = problem.stops, problem.tm
    n = len(stops)
    zones = np.array([o.zone for o in stops], dtype=int)
    serv = np.array([o.service for o in stops], dtype=float)
    early = np.array([o.earliest for o in stops], dtype=float)
    ret = problem.return_to_depot

    def h(cur: int, mask: int, clock: float) -> float:
        if kind == "zero":
            return 0.0
        un = [k for k in range(n) if not mask >> k & 1]
        if not un:
            return 0.0
        cur_zone = problem.depot if cur == DEPOT else stops[cur].zone
        LB = tm.M[tm._knot(clock)[0]] if kind == "naive" else tm.lower_bound_matrix(clock)
        uz = zones[un]
        # incoming into each unvisited k from cur or another unvisited stop
        inc = LB[np.append(uz, cur_zone)][:, uz]          # rows: sources (un..., cur)
        if len(un) > 1:
            sub = inc[:-1].copy()
            np.fill_diagonal(sub, np.inf)                 # u != k
            min_in = np.minimum(sub.min(axis=0), inc[-1])
        else:
            min_in = inc[-1]
        ret_lb = float(LB[uz, problem.depot].min()) if ret else 0.0
        if kind == "max_in":
            return float((min_in + serv[un]).max()) + ret_lb
        h_sum = float(min_in.sum() + serv[un].sum()) + ret_lb
        if kind == "naive":
            return h_sum
        h_wait = float((np.maximum(early[un], clock + min_in) + serv[un]).max()) - clock + ret_lb
        return max(h_sum, h_wait)

    return h


# --------------------------------------------------------------------------- best-first search
def solve_route(problem: RouteProblem, algo: str = "astar", heuristic: str = "admissible",
                hard: bool = True, trace: list | None = None, trace_limit: int = 40,
                max_expansions: int = 2_000_000) -> Route | None:
    """algo in {'ucs', 'greedy', 'astar'}. Returns the best Route or None (no feasible order)."""
    t0 = time.perf_counter()
    stops, tm, n = problem.stops, problem.tm, len(problem.stops)
    full = (1 << n) - 1
    hfun = make_heuristic(problem, "zero" if algo == "ucs" else heuristic)
    stats = {"algo": algo, "heuristic": heuristic if algo != "ucs" else "zero", "expanded": 0,
             "generated": 0, "pruned_infeasible": 0, "pruned_dominated": 0, "max_frontier": 0}

    # node: (cur, mask, clock, late, g, parent_id)
    nodes = []
    labels: dict = {}          # (cur, mask) -> list of (clock, late, node_id)
    alive = set()
    frontier = []
    tie = itertools.count()

    def push(cur, mask, clock, late, g, parent, terminal=False):
        key = (cur, mask, terminal)
        lab = labels.setdefault(key, [])
        for (c2, l2, _id) in lab:
            if c2 <= clock + 1e-9 and l2 <= late + 1e-9:   # an existing label dominates
                stats["pruned_dominated"] += 1
                return
        keep = []
        for (c2, l2, nid) in lab:
            if clock <= c2 + 1e-9 and late <= l2 + 1e-9:
                alive.discard(nid)
            else:
                keep.append((c2, l2, nid))
        nid = len(nodes)
        nodes.append((cur, mask, clock, late, g, parent, terminal))
        keep.append((clock, late, nid))
        labels[key] = keep
        alive.add(nid)
        hv = 0.0 if terminal else hfun(cur, mask, clock)
        pri = {"ucs": g, "astar": g + hv, "greedy": hv}[algo]
        heapq.heappush(frontier, (pri, g, next(tie), nid, hv))
        stats["generated"] += 1
        stats["max_frontier"] = max(stats["max_frontier"], len(frontier))

    def goal_transition(cur, mask, clock, late, g, nid):
        """All stops served: close the route (return leg, shift end)."""
        cz = problem.depot if cur == DEPOT else stops[cur].zone
        end = clock + (tm.tt(cz, problem.depot, clock) if problem.return_to_depot and n else 0.0)
        over = max(0.0, end - problem.shift_end)
        if hard and over > 1e-9:
            stats["pruned_infeasible"] += 1
            return
        push(cur, mask, end, late + over, end - problem.start + LATE_PENALTY * (late + over), nid,
             terminal=True)

    push(DEPOT, 0, problem.start, 0.0, 0.0, None)
    goal = None
    while frontier:
        pri, g, _, nid, hv = heapq.heappop(frontier)
        if nid not in alive:
            continue
        cur, mask, clock, late, g, parent, terminal = nodes[nid]
        if terminal:
            goal = nid
            break
        alive.discard(nid)
        stats["expanded"] += 1
        if trace is not None and len(trace) < trace_limit:
            trace.append({"step": stats["expanded"], "at": "depot" if cur == DEPOT else stops[cur].id,
                          "visited": [stops[k].id for k in range(n) if mask >> k & 1],
                          "clock": clock, "g": g, "h": hv, "f": g + hv})
        if stats["expanded"] > max_expansions:
            break
        if mask == full:
            goal_transition(cur, mask, clock, late, g, nid)
            continue
        cz = problem.depot if cur == DEPOT else stops[cur].zone
        for j in range(n):
            if mask >> j & 1:
                continue
            o = stops[j]
            a = clock + tm.tt(cz, o.zone, clock)
            b = max(a, o.earliest)
            lj = max(0.0, b - o.latest)
            if hard and lj > 1e-9:
                stats["pruned_infeasible"] += 1
                continue
            c2 = b + o.service
            push(j, mask | 1 << j, c2, late + lj, c2 - problem.start + LATE_PENALTY * (late + lj), nid)

    stats["time_s"] = time.perf_counter() - t0
    if goal is None:
        return None
    seq, k = [], nodes[goal][5]
    while k is not None and nodes[k][0] != DEPOT:
        seq.append(nodes[k][0])
        k = nodes[k][5]
    seq.reverse()
    route = simulate(problem, seq)
    route.stats = stats
    return route


def plan_van(problem: RouteProblem, algo: str = "astar", trace=None) -> Route:
    """Hard-window search first; if no on-time order exists, return the least-late route."""
    if not problem.stops:
        r = simulate(problem, [])
        r.stats = {"algo": algo, "expanded": 0, "generated": 0, "time_s": 0.0}
        return r
    r = solve_route(problem, algo=algo, hard=True, trace=trace)
    if r is not None:
        return r
    r = solve_route(problem, algo=algo, hard=False)
    r.stats["soft_fallback"] = True
    return r