"""The layout solver: items placed in atoms so that every size stays valid.

Used by a designer (tools/pagesets/solve.py) to propose where items live; the
result is written to a layout YAML, reviewed, and committed. The build never
runs the solver: it reads the committed positions and checks them.

Search
------
1. Construct: stage by stage, give each item of the stage an atom in a cell of
   that size no earlier item covers (depth-first, most-constrained item first,
   atoms in random order, honouring an item's allowed atoms). With as many
   items per stage as the size grows by, every cell of every size is filled.
2. Improve: simulated annealing over moves of one item to another atom, or
   swaps of two items' atoms, keeping every pair of items compatible (never
   in one cell at a size where both are present).
3. Polish: within each stage, the cheapest assignment of that stage's items
   to that stage's atoms (Hungarian algorithm) — swapping items of one stage
   can never make a map invalid.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass, field
from typing import Callable

from .grid import Ladder


def hungarian(cost: list[list[float]]) -> tuple[float, list[int]]:
    """Minimum-cost perfect matching of a square matrix: (total, col for each row)."""
    n = len(cost)
    if n == 0:
        return 0.0, []
    INF = float("inf")
    u = [0.0] * (n + 1)
    v = [0.0] * (n + 1)
    p = [0] * (n + 1)
    way = [0] * (n + 1)
    for i in range(1, n + 1):
        p[0] = i
        j0 = 0
        minv = [INF] * (n + 1)
        used = [False] * (n + 1)
        while True:
            used[j0] = True
            i0 = p[j0]
            delta = INF
            j1 = 0
            for j in range(1, n + 1):
                if not used[j]:
                    cur = cost[i0 - 1][j - 1] - u[i0] - v[j]
                    if cur < minv[j]:
                        minv[j] = cur
                        way[j] = j0
                    if minv[j] < delta:
                        delta = minv[j]
                        j1 = j
            for j in range(n + 1):
                if used[j]:
                    u[p[j]] += delta
                    v[j] -= delta
                else:
                    minv[j] -= delta
            j0 = j1
            if p[j0] == 0:
                break
        while True:
            j1 = way[j0]
            p[j0] = p[j1]
            j0 = j1
            if j0 == 0:
                break
    assign = [0] * n
    for j in range(1, n + 1):
        if p[j]:
            assign[p[j] - 1] = j - 1
    return sum(cost[i][assign[i]] for i in range(n)), assign


@dataclass
class Item:
    key: str
    stage: int  # ladder index at which it first appears
    pinned: int | None = None  # atom, when fixed by the designer
    allowed: set[int] | None = None  # atoms it may take (None: any)
    data: dict = field(default_factory=dict)


class StageSolver:
    def __init__(
        self,
        ladder: Ladder,
        items: list[Item],
        cost: Callable[[Item, int], float],
        full: bool = True,
        seed: int = 1,
    ):
        self.ladder = ladder
        self.items = items
        self.cost = cost
        self.full = full
        self.rng = random.Random(seed)
        self.k = len(ladder.sizes)
        self.by_stage: list[list[Item]] = [[] for _ in range(self.k)]
        for it in items:
            self.by_stage[it.stage].append(it)
            if it.pinned is not None:
                it.allowed = {it.pinned}
        if full:
            steps = ladder.capacity_steps()
            counts = [len(s) for s in self.by_stage]
            if counts != steps:
                raise ValueError(f"items per stage {counts} must equal capacity steps {steps}")
        self._cc: dict[tuple[str, int], float] = {}

    def c(self, it: Item, atom: int) -> float:
        key = (it.key, atom)
        if key not in self._cc:
            self._cc[key] = self.cost(it, atom)
        return self._cc[key]

    def _fits(self, atom: int, stage: int, placed: dict[str, tuple[int, int]], skip: tuple[str, ...] = ()) -> bool:
        lad = self.ladder
        for key, (a, s) in placed.items():
            if key in skip:
                continue
            if not lad.compatible(atom, stage, a, s):
                return False
        return True

    # ── 1. construct ──────────────────────────────────────────────────────
    def construct(self, tries: int = 400, node_limit: int = 20000) -> dict[str, int]:
        lad = self.ladder
        for _ in range(tries):
            placed: dict[str, tuple[int, int]] = {}
            ok = True
            for k in range(self.k):
                items = self.by_stage[k]
                taken = {a for a, _ in placed.values()}
                covered = {lad.block[k][a] for a, s in placed.values()}
                opts = {}
                for it in items:
                    cand = it.allowed if it.allowed is not None else range(lad.n)
                    o = [a for a in cand if a not in taken and lad.block[k][a] not in covered and self._fits(a, k, placed)]
                    self.rng.shuffle(o)
                    # prefer cheaper atoms, with some randomness
                    o.sort(key=lambda a, it=it: self.c(it, a) * (0.7 + 0.6 * self.rng.random()))
                    opts[it.key] = o
                order = sorted(items, key=lambda it: len(opts[it.key]))
                nodes = [0]
                chosen: dict[str, int] = {}

                def dfs(i: int) -> bool:
                    if i == len(order):
                        return True
                    nodes[0] += 1
                    if nodes[0] > node_limit:
                        return False
                    it = order[i]
                    blocks = {lad.block[k][a] for a in chosen.values()}
                    for a in opts[it.key]:
                        if lad.block[k][a] in blocks:
                            continue
                        if any(not lad.compatible(a, k, b, k) for b in chosen.values()):
                            continue
                        chosen[it.key] = a
                        if dfs(i + 1):
                            return True
                        del chosen[it.key]
                    return False

                if not dfs(0):
                    ok = False
                    break
                for key, a in chosen.items():
                    placed[key] = (a, k)
            if ok:
                return {key: a for key, (a, _) in placed.items()}
        raise RuntimeError("no valid placement found: loosen the anchors or try another seed")

    # ── 2. improve ────────────────────────────────────────────────────────
    def total(self, where: dict[str, int]) -> float:
        return sum(self.c(it, where[it.key]) for it in self.items)

    def anneal(self, where: dict[str, int], iters: int) -> dict[str, int]:
        lad = self.ladder
        stage = {it.key: it.stage for it in self.items}
        item = {it.key: it for it in self.items}
        movable = [it for it in self.items if it.pinned is None]
        placed = {k: (a, stage[k]) for k, a in where.items()}
        occupied = {a: k for k, a in where.items()}
        cur = self.total(where)
        t0 = max(cur * 0.01, 1e-6)
        t1 = t0 * 1e-4
        best, best_cost = dict(where), cur
        for step in range(iters):
            t = t0 * (t1 / t0) ** (step / max(1, iters - 1))
            it = self.rng.choice(movable)
            a_old = placed[it.key][0]
            cand = list(it.allowed) if it.allowed is not None else None
            b = self.rng.choice(cand) if cand else self.rng.randrange(lad.n)
            if b == a_old:
                continue
            other = occupied.get(b)
            if other is None:
                if not self._fits(b, stage[it.key], placed, skip=(it.key,)):
                    continue
                delta = self.c(it, b) - self.c(it, a_old)
                if delta <= 0 or self.rng.random() < math.exp(-delta / t):
                    placed[it.key] = (b, stage[it.key])
                    del occupied[a_old]
                    occupied[b] = it.key
                    cur += delta
            else:
                ot = item[other]
                if ot.pinned is not None or (ot.allowed is not None and a_old not in ot.allowed):
                    continue
                trial = dict(placed)
                trial[it.key] = (b, stage[it.key])
                trial[other] = (a_old, stage[other])
                if not self._fits(b, stage[it.key], trial, skip=(it.key,)) or not self._fits(a_old, stage[other], trial, skip=(other,)):
                    continue
                delta = self.c(it, b) + self.c(ot, a_old) - self.c(it, a_old) - self.c(ot, b)
                if delta <= 0 or self.rng.random() < math.exp(-delta / t):
                    placed = trial
                    occupied[b], occupied[a_old] = it.key, other
                    cur += delta
            if cur < best_cost - 1e-12:
                best, best_cost = {k: a for k, (a, _) in placed.items()}, cur
        return best

    # ── 3. polish ─────────────────────────────────────────────────────────
    def polish(self, where: dict[str, int]) -> dict[str, int]:
        out = dict(where)
        for k in range(self.k):
            items = [it for it in self.by_stage[k] if it.pinned is None]
            atoms = [out[it.key] for it in items]
            if len(items) < 2:
                continue
            big = 1e9
            m = [[self.c(it, a) if (it.allowed is None or a in it.allowed) else big for a in atoms] for it in items]
            _, assign = hungarian(m)
            for r, it in enumerate(items):
                out[it.key] = atoms[assign[r]]
        return out

    def solve(self, restarts: int = 6, iters: int = 20000) -> tuple[dict[str, int], float]:
        best, best_cost = None, float("inf")
        for _ in range(restarts):
            w = self.polish(self.anneal(self.construct(), iters))
            c = self.total(w)
            if c < best_cost:
                best, best_cost = w, c
        assert best is not None
        return best, best_cost