"""Order -> van assignment as a CSP.

Variables   X_o for every order o
Domains     the vans {0, 1, 2}
Constraints
  unary      window reachable from the depot by that van: leaving at shift start,
             max(shift_start + T(depot, o), e_o) <= l_o, and (if returning) the van can still
             get back to the depot before its shift ends.
  binary     X_a != X_b  ("different van") when the predicted times make it impossible to
             serve both windows in EITHER order (depot -> a -> b and depot -> b -> a), using the
             earliest shift start (a necessary condition for every van). These edges come from
             the NN time matrix. The integration loop adds more ("learned") edges.
  global     capacity: sum of demands <= van capacity, number of stops <= max_stops.

Solvers (switchable parts)
  BT        chronological backtracking, checks constraints against past assignments
  FC        forward checking: after X_i = v, delete v from unassigned neighbours' domains and
            from any order that no longer fits in v's remaining capacity / stop budget
  AC-3      arc consistency on the "different van" constraints, run once before search and
            then after every assignment (MAC) on top of FC
Variable order: static, most-constrained (highest degree) first, ties by earliest window.
Value order: least-loaded van first (deterministic; identical for all three solvers, so they
return the same first solution and differ only in how much of the tree they explore).

Counters: assignments (values given to variables), backtracks (assignments undone because
their subtree had no solution), checks, fc_prunes, ac3_revisions, ac3_prunes.
"""
from __future__ import annotations

import time
from collections import deque
from dataclasses import dataclass, field


def fmt_t(m: float) -> str:
    m = int(round(m))
    return f"{m // 60:02d}:{m % 60:02d}"


@dataclass
class CSPResult:
    status: str                    # 'solved' | 'unsat' | 'limit'
    assignment: dict               # order id -> van id
    stats: dict = field(default_factory=dict)


class AssignmentCSP:
    def __init__(self, instance, tm, learned_edges=()):
        self.inst, self.tm = instance, tm
        self.orders = list(instance.orders)
        self.ids = [o.id for o in self.orders]
        self.vans = list(instance.vans)
        self.unary_reasons: dict = {}
        self.domains0 = {o.id: self._unary(o) for o in self.orders}
        self.edges: dict = {}                       # frozenset({a,b}) -> 'nn' | 'learned'
        self.edge_reason: dict = {}
        self._pairwise()
        for a, b in learned_edges:
            self.add_edge(a, b, "learned")

    # ------------------------------------------------------------------ constraint building
    def _unary(self, o):
        dom, depot, tm = [], self.inst.depot, self.tm
        for v in self.vans:
            t = v.shift_start
            arr = t + tm.tt(depot, o.zone, t)
            beg = max(arr, o.earliest)
            back = beg + o.service + (tm.tt(o.zone, depot, beg + o.service) if self.inst.return_to_depot else 0)
            if beg > o.latest + 1e-9:
                self.unary_reasons.setdefault(o.id, []).append(
                    f"van {v.id}: earliest arrival {fmt_t(arr)} is after the window closes {fmt_t(o.latest)}")
            elif back > v.shift_end + 1e-9:
                self.unary_reasons.setdefault(o.id, []).append(
                    f"van {v.id}: cannot be back at the depot before shift end {fmt_t(v.shift_end)}")
            elif o.demand > v.capacity:
                self.unary_reasons.setdefault(o.id, []).append(
                    f"van {v.id}: demand {o.demand} exceeds capacity {v.capacity}")
            else:
                dom.append(v.id)
        return dom

    def pair_feasible(self, a, b) -> bool:
        """Can one van serve a then b on time, starting from the depot at the earliest shift?"""
        tm, depot = self.tm, self.inst.depot
        s0 = min(v.shift_start for v in self.vans)
        ta = max(s0 + tm.tt(depot, a.zone, s0), a.earliest)
        if ta > a.latest + 1e-9:
            return False
        d = ta + a.service
        tb = max(d + tm.tt(a.zone, b.zone, d), b.earliest)
        return tb <= b.latest + 1e-9

    def _pairwise(self):
        for i, a in enumerate(self.orders):
            for b in self.orders[i + 1:]:
                if not self.pair_feasible(a, b) and not self.pair_feasible(b, a):
                    self.add_edge(a.id, b.id, "nn")

    def add_edge(self, a, b, source):
        key = frozenset((a, b))
        if a != b and key not in self.edges:
            self.edges[key] = source

    def neighbors(self):
        nb = {i: set() for i in self.ids}
        for e in self.edges:
            a, b = tuple(e)
            nb[a].add(b); nb[b].add(a)
        return nb

    # ------------------------------------------------------------------ solving
    def solve(self, fc: bool = False, ac3: bool = False, max_assignments: int = 300_000,
              time_limit: float = 60.0, var_order: str = "degree") -> CSPResult:
        """var_order: 'degree' (most constrained first) or 'input' (order ids as given)."""
        t0 = time.perf_counter()
        nb = self.neighbors()
        demand = {o.id: o.demand for o in self.orders}
        cap = {v.id: v.capacity for v in self.vans}
        maxs = {v.id: v.max_stops for v in self.vans}
        D = {i: set(self.domains0[i]) for i in self.ids}
        st = {"solver": "BT" + ("+FC" if fc else "") + ("+AC3" if ac3 else ""), "assignments": 0,
              "backtracks": 0, "checks": 0, "fc_prunes": 0, "ac3_revisions": 0, "ac3_prunes": 0}
        if var_order == "degree":
            order = sorted(self.ids, key=lambda i: (-len(nb[i]), self.inst.order(i).earliest, i))
        else:
            order = list(self.ids)
        A: dict = {}
        load = {v: 0 for v in cap}
        count = {v: 0 for v in cap}
        trail: list = []                    # (var, value) removed from domains

        def remove(var, val):
            D[var].discard(val)
            trail.append((var, val))

        def undo(mark):
            while len(trail) > mark:
                var, val = trail.pop()
                D[var].add(val)

        def ac3_run(queue) -> bool:
            q = deque(queue)
            while q:
                x, y = q.popleft()
                st["ac3_revisions"] += 1
                if len(D[y]) == 1:          # x != y : only a singleton domain can prune
                    (b,) = tuple(D[y])
                    if b in D[x]:
                        remove(x, b); st["ac3_prunes"] += 1
                        if not D[x]:
                            return False
                        q.extend((z, x) for z in nb[x] if z != y and z not in A)
            return True

        if any(not D[i] for i in self.ids):
            st["time_s"] = time.perf_counter() - t0
            return CSPResult("unsat", {}, st)
        if ac3 and not ac3_run([(x, y) for x in self.ids for y in nb[x]]):
            st["time_s"] = time.perf_counter() - t0
            return CSPResult("unsat", {}, st)

        status = {"limit": False}

        def consistent(i, v) -> bool:
            st["checks"] += 1
            if load[v] + demand[i] > cap[v] or count[v] + 1 > maxs[v]:
                return False
            return all(A.get(j) != v for j in nb[i])

        def bt(k) -> bool:
            if k == len(order):
                return True
            if st["assignments"] >= max_assignments or time.perf_counter() - t0 > time_limit:
                status["limit"] = True
                return False
            i = order[k]
            for v in sorted(D[i], key=lambda v: (count[v], v)):
                if not fc and not consistent(i, v):
                    continue
                A[i] = v; load[v] += demand[i]; count[v] += 1
                st["assignments"] += 1
                mark = len(trail)
                ok = True
                if fc:
                    changed = []
                    for u in order[k + 1:]:
                        if v in D[u] and (u in nb[i] or load[v] + demand[u] > cap[v] or count[v] >= maxs[v]):
                            remove(u, v); st["fc_prunes"] += 1; changed.append(u)
                            if not D[u]:
                                ok = False
                                break
                    if ok and ac3:
                        ok = ac3_run([(z, u) for u in changed for z in nb[u] if z not in A])
                if ok and bt(k + 1):
                    return True
                undo(mark)
                del A[i]; load[v] -= demand[i]; count[v] -= 1
                st["backtracks"] += 1
                if status["limit"]:
                    return False
            return False

        found = bt(0)
        st["time_s"] = time.perf_counter() - t0
        if found:
            return CSPResult("solved", dict(A), st)
        return CSPResult("limit" if status["limit"] else "unsat", {}, st)

    # ------------------------------------------------------------------ diagnostics
    def max_clique(self):
        nb = self.neighbors()
        best: list = []

        def bk(R, P, X):
            nonlocal best
            if not P and not X:
                if len(R) > len(best):
                    best = list(R)
                return
            if len(R) + len(P) <= len(best):
                return
            u = max(P | X, key=lambda z: len(nb[z] & P))
            for w in list(P - nb[u]):
                bk(R | {w}, P & nb[w], X & nb[w])
                P = P - {w}; X = X | {w}

        bk(set(), set(self.ids), set())
        return sorted(best)


def check_assignment(csp: AssignmentCSP, assignment: dict) -> list:
    """Returns a list of violated constraints (empty = valid). Used by the tests."""
    bad = []
    for i in csp.ids:
        if assignment.get(i) not in csp.domains0[i]:
            bad.append(f"order {i}: van {assignment.get(i)} not in domain {csp.domains0[i]}")
    for e in csp.edges:
        a, b = tuple(e)
        if assignment.get(a) == assignment.get(b):
            bad.append(f"orders {a},{b} share van {assignment.get(a)} but must differ")
    for v in csp.vans:
        mine = [csp.inst.order(i) for i in csp.ids if assignment.get(i) == v.id]
        if sum(o.demand for o in mine) > v.capacity:
            bad.append(f"van {v.id}: capacity exceeded")
        if len(mine) > v.max_stops:
            bad.append(f"van {v.id}: too many stops")
    return bad