"""Set-partitioning matheuristic (--solver setpart): route pool + MILP.

The airline crew-scheduling pattern at field-adjuster scale. Heuristics
are excellent at building individual routes but must commit to ONE way
of splitting claims across adjusters; the remaining gap on big days
lives in that split. So:

1. Harvest routes from several diverse heuristic runs (OR-Tools GLS,
   the hybrid, PyVRP under different random seeds).
2. Enrich the pool: try every harvested route on every other eligible
   adjuster (skills/territory/shift permitting), add single-claim
   routes, and re-sequence every candidate to per-route optimality by
   exhaustive branch and bound (routes are short, so this is exact and
   fast).
3. A tiny MILP - one binary per candidate route - picks the best
   combination: each adjuster drives at most one route, each claim is
   covered exactly once or dropped at its usual penalty.

The selection MILP stays small (hundreds of binaries) no matter how
large the day is, because it only chooses among precomputed routes -
it never sequences anything. That is what makes this scale where the
full exact MILP (milp_solver.py) hits its ~25-claim wall.

Guarantee: every source solution is itself a feasible selection in the
pool, so the chosen schedule is never worse than the best heuristic
run that fed it.
"""

from __future__ import annotations

import time

import pulp

import config
import solver as ortools_solver
from data_gen import Adjuster, Claim
from milp_solver import _pick_engine
from solver import Route, Solution, Stop


def _latest_start(c: Claim) -> int:
    """Latest service start: the visit must also finish by window end."""
    return max(c.window_start, c.window_end - c.service_minutes)


def _simulate(seq, k, adj, node, by_id, travel_min):
    """Forward earliest-time pass. Returns total travel minutes, or None
    if infeasible. First-stop wait is unlimited (absorbed by delaying the
    home departure, as OR-Tools does); later waits obey the slack cap."""
    t, prev, travel = adj.shift_start, k, 0
    for pos, cid in enumerate(seq):
        c = by_id[cid]
        leg = travel_min[prev][node[cid]]
        service_start = max(c.window_start, t + leg)
        if service_start > _latest_start(c):
            return None
        if pos > 0 and service_start - (t + leg) > config.MAX_WAIT_MINUTES:
            return None
        travel += leg
        t = service_start + c.service_minutes
        prev = node[cid]
    travel += travel_min[prev][k]
    if t + travel_min[prev][k] > adj.shift_end:
        return None
    return travel


def _optimal_sequence(cids, k, adj, node, by_id, travel_min,
                      budget_s: float = 5.0):
    """Exhaustive branch and bound over stop orders for one claim set.
    Returns (travel_minutes, best_order) or (None, None) if no feasible
    order exists. Wait rule mirrors _simulate (conservative, never
    accepts what the OR-Tools model would reject)."""
    deadline = time.time() + budget_s
    pts = [k] + [node[c] for c in cids]
    min_in = {c: min(travel_min[p][node[c]] for p in pts if p != node[c])
              for c in cids}
    min_home = {c: travel_min[node[c]][k] for c in cids}
    best = [None, None]   # travel, order
    seq = []

    def dfs(prev, t, remaining, travel):
        if time.time() > deadline:
            return
        if not remaining:
            total = travel + travel_min[prev][k]
            if ((best[0] is None or total < best[0])
                    and t + travel_min[prev][k] <= adj.shift_end):
                best[0], best[1] = total, list(seq)
            return
        lb = travel + sum(min_in[c] for c in remaining) \
            + min(min_home[c] for c in remaining)
        if best[0] is not None and lb >= best[0]:
            return
        for cid in list(remaining):
            c = by_id[cid]
            leg = travel_min[prev][node[cid]]
            service_start = max(c.window_start, t + leg)
            if service_start > _latest_start(c):
                continue
            if seq and service_start - (t + leg) > config.MAX_WAIT_MINUTES:
                continue
            remaining.remove(cid)
            seq.append(cid)
            dfs(node[cid], service_start + c.service_minutes,
                remaining, travel + leg)
            seq.pop()
            remaining.add(cid)

    dfs(k, adj.shift_start, set(cids), 0)
    return best[0], best[1]


def _eligible_for_all(adj, k, cids, node, by_id, miles) -> bool:
    return all(config.is_eligible(adj, by_id[cid], miles[k][node[cid]])
               for cid in cids)


def _build_solution(chosen, adjusters, claims, miles, travel_min):
    """Chosen = {adjuster_id: [claim_ids in order]}. Same earliest-time
    reconstruction as milp_solver.solve_as_backend."""
    n_adj = len(adjusters)
    node = {c.claim_id: n_adj + i for i, c in enumerate(claims)}
    by_id = {c.claim_id: c for c in claims}
    routes, served = [], set()
    for k, a in enumerate(adjusters):
        seq = chosen.get(a.adjuster_id, [])
        route = Route(adjuster=a)
        prev, t = k, a.shift_start
        for cid in seq:
            c = by_id[cid]
            leg = travel_min[prev][node[cid]]
            arrival = max(c.window_start, t + leg)
            route.stops.append(Stop(
                claim=c, arrival_min=arrival,
                departure_min=arrival + c.service_minutes,
                travel_miles_from_prev=miles[prev][node[cid]],
                travel_min_from_prev=leg))
            route.total_miles += miles[prev][node[cid]]
            route.total_travel_min += leg
            route.total_service_min += c.service_minutes
            served.add(cid)
            prev, t = node[cid], arrival + c.service_minutes
        if seq:
            route.start_min = (route.stops[0].arrival_min
                               - route.stops[0].travel_min_from_prev)
            route.total_miles += miles[prev][k]
            route.total_travel_min += travel_min[prev][k]
            route.end_min = t + travel_min[prev][k]
        else:
            route.start_min = route.end_min = a.shift_start
        routes.append(route)

    dropped = [c for c in claims if c.claim_id not in served]
    unservable = [c for i, c in enumerate(claims)
                  if not any(config.is_eligible(a, c, miles[k][n_adj + i])
                             for k, a in enumerate(adjusters))]
    travel_total = sum(r.total_travel_min for r in routes)
    return Solution(
        routes=routes, dropped=dropped, unservable=unservable,
        objective=travel_total + sum(
            config.effective_penalty(c.priority, c.age_days)
            for c in dropped),
        total_miles=sum(r.total_miles for r in routes),
        total_travel_min=travel_total)


def solve_setpartition(adjusters: list[Adjuster], claims: list[Claim],
                       miles: list[list[float]],
                       travel_min: list[list[int]],
                       time_limit_s: int = 60, engine: str = "auto",
                       lunch_break: bool = False,
                       balance: bool = False,
                       ) -> tuple[Solution | None, dict]:
    """Run the pool-then-select matheuristic. Returns (Solution, info).

    The time limit is split across the pool-generating heuristic runs;
    pool enrichment and the selection MILP add a few seconds on top.

    Toggles follow the airline recipe: the rules live inside the
    columns, never in the selection model. Sources run toggled where
    they can (ortools, hybrid; PyVRP seeds get balance natively), then
    EVERY candidate card is legalized and priced: with lunch on, a
    card that cannot host a 30-minute break (proven by the same
    interval machinery as the sequence backends) is ejected from the
    pool; with balance on, each card's cost gains BALANCE_COEFFICIENT
    times its exact minimal span. The harvested source orders are
    always offered alongside their re-sequenced variants, so a toggled
    source plan remains a feasible selection and the never-worse
    guarantee holds against toggled sources.
    """
    if any(c.extra_windows for c in claims):
        raise ValueError(
            "split-availability claims (multiple time windows) are not "
            "supported by this backend yet - use ortools, cpsat, or "
            "milp (exact)")
    n_adj = len(adjusters)
    node = {c.claim_id: n_adj + i for i, c in enumerate(claims)}
    by_id = {c.claim_id: c for c in claims}
    t_start = time.time()

    # ---- 1. Diverse heuristic runs -----------------------------------
    T = max(10, time_limit_s)
    budgets = {"ortools": max(3, int(0.20 * T)),
               "hybrid": max(4, int(0.35 * T)),
               "pyvrp-s1": max(3, int(0.125 * T)),
               "pyvrp-s2": max(3, int(0.125 * T))}
    sources: dict[str, Solution] = {}
    sol = ortools_solver.solve(adjusters, claims, miles, travel_min,
                               time_limit_s=budgets["ortools"],
                               lunch_break=lunch_break, balance=balance)
    if sol is not None:
        sources["ortools"] = sol
    try:
        import pyvrp_solver
        sol = pyvrp_solver.solve_hybrid(adjusters, claims, miles,
                                        travel_min,
                                        time_limit_s=budgets["hybrid"],
                                        lunch_break=lunch_break,
                                        balance=balance)
        if sol is not None:
            sources["hybrid"] = sol
        for name, seed in (("pyvrp-s1", 1), ("pyvrp-s2", 2)):
            sol = pyvrp_solver.solve(adjusters, claims, miles, travel_min,
                                     time_limit_s=budgets[name], seed=seed,
                                     balance=balance)
            if sol is not None:
                sources[name] = sol
    except ImportError:
        pass
    if not sources:
        return None, {"status": "no heuristic solution"}
    # The never-worse yardstick must obey the same rules as the
    # selection: with lunch on, the PyVRP seed runs are break-blind,
    # so their objectives are not comparable - they still donate
    # cards (each individually legalized), but cannot set the floor.
    yardstick = {name: s for name, s in sources.items()
                 if not lunch_break or name in ("ortools", "hybrid")}
    if not yardstick:
        yardstick = sources
    best_source_name = min(yardstick,
                           key=lambda s: yardstick[s].objective)
    best_source = yardstick[best_source_name]

    # ---- 2. Route pool: harvest, transplant, re-sequence -------------
    # key = (adjuster_id, frozenset of claim ids) -> (cost, order); the
    # cost is travel minutes, plus the card's span cost under balance.
    pool: dict[tuple, tuple[int, list[str]]] = {}
    k_of = {a.adjuster_id: k for k, a in enumerate(adjusters)}
    adj_of = {a.adjuster_id: a for a in adjusters}
    ejected = [0]
    if lunch_break or balance:
        import sequencer as _seq

    def offer(adj_id, cost, order):
        if cost is None:
            return
        if lunch_break or balance:
            timed = _seq._min_span_times(list(order), k_of[adj_id],
                                         adj_of[adj_id], node, by_id,
                                         travel_min,
                                         lunch_break=lunch_break)
            if timed is None:            # card cannot host the break
                ejected[0] += 1
                return
            if balance:
                cost = cost + config.BALANCE_COEFFICIENT * (
                    timed[2] - timed[1])
        key = (adj_id, frozenset(order))
        if key not in pool or cost < pool[key][0]:
            pool[key] = (cost, list(order))

    harvested = []
    for src in sources.values():
        for r in src.routes:
            if r.stops:
                harvested.append((r.adjuster.adjuster_id,
                                  [s.claim.claim_id for s in r.stops]))

    for src_adj_id, order in harvested:
        for k, a in enumerate(adjusters):
            if not _eligible_for_all(a, k, order, node, by_id, miles):
                continue
            # the harvested order itself is always offered, so a
            # toggled source plan stays selectable card-for-card
            offer(a.adjuster_id,
                  _simulate(order, k, a, node, by_id, travel_min), order)
            if len(order) <= 11:
                cost, best_order = _optimal_sequence(
                    order, k, a, node, by_id, travel_min)
                if cost is not None:
                    offer(a.adjuster_id, cost, best_order)

    # Leave-one-out sub-routes: give the selection MILP fine-grained
    # pieces to recombine (whole routes from near-optimal runs overlap
    # too much to mix). Each subset is re-sequenced optimally.
    for (adj_id, cset), (_, order) in list(pool.items()):
        if len(order) < 2 or len(order) > 11:
            continue
        a, k = adj_of[adj_id], k_of[adj_id]
        for drop_cid in order:
            sub = [cid for cid in order if cid != drop_cid]
            cost, best_order = _optimal_sequence(sub, k, a, node, by_id,
                                                 travel_min)
            if cost is not None:
                offer(adj_id, cost, best_order)

    for c in claims:                      # single-claim fallback columns
        for k, a in enumerate(adjusters):
            i = node[c.claim_id]
            if config.is_eligible(a, c, miles[k][i]):
                cost = _simulate([c.claim_id], k, a, node, by_id,
                                 travel_min)
                offer(a.adjuster_id, cost, [c.claim_id])

    # ---- 3. Selection MILP -------------------------------------------
    engine = _pick_engine(engine)
    prob = pulp.LpProblem("route_selection", pulp.LpMinimize)
    y = {key: pulp.LpVariable(f"y_{i}", cat="Binary")
         for i, key in enumerate(pool)}
    d = {c.claim_id: pulp.LpVariable(f"d_{c.claim_id}", cat="Binary")
         for c in claims}
    prob += (pulp.lpSum(pool[key][0] * var for key, var in y.items())
             + pulp.lpSum(config.effective_penalty(c.priority, c.age_days)
                          * d[c.claim_id] for c in claims))
    for a in adjusters:
        cols = [var for key, var in y.items() if key[0] == a.adjuster_id]
        if cols:
            prob += pulp.lpSum(cols) <= 1
    for c in claims:
        covering = [var for key, var in y.items() if c.claim_id in key[1]]
        prob += pulp.lpSum(covering) + d[c.claim_id] == 1

    mip_t0 = time.time()
    if engine == "gurobi":
        try:
            prob.solve(pulp.GUROBI(msg=False, timeLimit=30, gapRel=0.0))
        except Exception:
            engine = "highs"
    if engine == "highs":
        prob.solve(pulp.HiGHS(msg=False, timeLimit=30, gapRel=0.0))
    proven = prob.sol_status == pulp.LpSolutionOptimal

    chosen = {}
    for key, var in y.items():
        if var.value() is not None and var.value() > 0.5:
            chosen[key[0]] = pool[key][1]
    sp_sol = _build_solution(chosen, adjusters, claims, miles, travel_min)
    if lunch_break or balance:
        # restamp times/endpoints with the lunch-and-span-aware timing
        # and report the same objective the toggled backends report
        for k, route in enumerate(sp_sol.routes):
            if not route.stops:
                continue
            order = [s.claim.claim_id for s in route.stops]
            timed = _seq._min_span_times(order, k, adjusters[k], node,
                                         by_id, travel_min,
                                         lunch_break=lunch_break)
            times, t_start, t_end = timed
            for stop, t in zip(route.stops, times):
                stop.arrival_min = t
                stop.departure_min = t + stop.claim.service_minutes
            route.start_min, route.end_min = t_start, t_end
        if balance:
            sp_sol.objective += config.BALANCE_COEFFICIENT * sum(
                r.end_min - r.start_min for r in sp_sol.routes
                if r.stops)

    # Safety net: the pool contains every source solution, so the MILP
    # can never legitimately do worse than the best of them.
    if sp_sol.objective > best_source.objective:
        sp_sol = best_source

    info = {
        "engine": engine,
        "selection_proven_optimal": proven,
        "pool_columns": len(pool),
        "ejected_cards": ejected[0],
        "toggles": {"lunch_break": lunch_break, "balance": balance},
        "sources": {name: s.objective for name, s in sources.items()},
        "best_source": best_source_name,
        "best_source_objective": best_source.objective,
        "objective": sp_sol.objective,
        "improvement_vs_best_source":
            best_source.objective - sp_sol.objective,
        "mip_seconds": round(time.time() - mip_t0, 2),
        "total_seconds": round(time.time() - t_start, 1),
    }
    return sp_sol, info