"""Layouts of the two sets: solving (designer), loading and checking (build)."""

from __future__ import annotations

import json
import math
from dataclasses import dataclass, field

from . import costs
from .data import HERE, Data, load_yaml
from .grid import Size, check_stages, fmt_point, ladder_for, place
from .solve import Item, StageSolver

BACK = "Back"  # the Back cell of a page (Growing Core: top-left; Fixed Paths: where its key was)


# ════════════════════════════════════════════════════════════════════════
# Growing Core
# ════════════════════════════════════════════════════════════════════════


@dataclass
class Placed:
    key: str
    kind: str  # word | page | back | modifier
    stage: int  # index into its ladder
    point: tuple[float, float]
    at: dict[str, tuple[int, int]] = field(default_factory=dict)


@dataclass
class CorePage:
    name: str
    label: str
    role: str
    sizes: list[Size]  # its own ladder (from its opener's stage)
    items: list[Placed]
    modifiers: dict[str, str]  # label -> modifier name


@dataclass
class CoreLayout:
    sizes: list[Size]
    core: list[Placed]
    pages: dict[str, CorePage]


def _core_design(data: Data):
    d = data.growing_core
    sizes = [Size.parse(s) for s in d["ladder"]]
    names = [s.name for s in sizes]
    core_items = []
    for stage_name, keys in d["core"].items():
        k = names.index(stage_name)
        for key in keys:
            kind = "page" if key in d["pages"] else "word"
            if kind == "word":
                data.entry(key)
            core_items.append((key, kind, k))
    return d, sizes, core_items


def _page_stages(page_sizes: list[Size], n_words: int, reserved: int = 1) -> list[int]:
    """Stage (index into page_sizes) of each word in priority order."""
    stages, k, used = [], 0, 0
    for _ in range(n_words):
        while k < len(page_sizes) and used >= page_sizes[k].cells - reserved:
            k += 1
        if k == len(page_sizes):
            raise SystemExit(f"a page has more words than its largest size holds ({page_sizes[-1]})")
        stages.append(k)
        used += 1
    return stages


def core_anchors(d: dict) -> dict[str, list[tuple[Size, tuple[int, int]]]]:
    """The designed first boards: word -> [(size, cell)]."""
    out: dict[str, list[tuple[Size, tuple[int, int]]]] = {}
    for name, text in (d.get("anchors") or {}).items():
        size = Size.parse(name)
        rows = [r for r in text.strip("\n").splitlines() if r.strip()]
        if len(rows) != size.rows:
            raise SystemExit(f"anchors {name}: {len(rows)} rows")
        for r, row in enumerate(rows):
            cells = [c.strip() for c in row.split("|")]
            if len(cells) != size.cols:
                raise SystemExit(f"anchors {name} row {r}: {len(cells)} cells")
            for c, key in enumerate(cells):
                out.setdefault(key, []).append((size, (r, c)))
    return out


def solve_core(data: Data, seed: int, restarts: int, iters: int) -> str:
    d, sizes, core_items = _core_design(data)
    ladder = ladder_for(tuple(s.name for s in sizes))
    regions = d["regions"]
    anchors = core_anchors(d)

    def role_of(key, kind):
        return "nav" if kind == "page" else data.entry(key).role

    items = []
    for key, kind, k in core_items:
        allowed = None
        if key in anchors:
            allowed = {a for a in range(ladder.n) if all(place(ladder.points[a], s) == cell for s, cell in anchors[key])}
            if not allowed:
                raise SystemExit(f"anchors: no point satisfies every designed board for {key!r}")
        items.append(Item(key, k, allowed=allowed, data={"role": role_of(key, kind), "n": len(sizes) - k}))

    def cost(it: Item, atom: int) -> float:
        point = ladder.points[atom]
        v, u = point
        reg = regions[it.data["role"]]
        c = (u - reg["u"]) ** 2 + (0.3 * (v - reg["v"]) ** 2 if "v" in reg else 0.0)
        c *= math.sqrt(it.data["n"])
        return c

    solver = StageSolver(ladder, items, cost, full=True, seed=seed)
    assign, total = solver.solve(restarts=restarts, iters=iters)
    core = [Placed(key, kind, k, ladder.points[assign[key]]) for key, kind, k in core_items]
    missed = [
        f"{p.key} at {size}" for p in core for size, cell in anchors.get(p.key, []) if place(p.point, size) != cell
    ]
    if missed:
        raise SystemExit("the solver could not keep these anchors (try more --restarts): " + ", ".join(missed))

    pages = {}
    opener_stage = {key: k for key, kind, k in core_items if kind == "page"}
    for name, spec in d["pages"].items():
        psizes = sizes[opener_stage[name]:]
        plad = ladder_for(tuple(s.name for s in psizes))
        words = [("modifier", m["label"]) for m in spec.get("modifiers", [])] + [("word", w) for w in spec["words"]]
        for kind, w in words:
            if kind == "word":
                data.entry(w)
        stages = _page_stages(psizes, len(words))
        pitems = [Item(BACK, 0, pinned=0)] + [
            Item(w, s, data={"rank": i}) for i, ((kind, w), s) in enumerate(zip(words, stages))
        ]

        def pcost(it: Item, atom: int, plad=plad) -> float:
            if it.key == BACK:
                return 0.0
            v, u = plad.points[atom]
            return (v + 0.6 * u) / math.sqrt(it.data["rank"] + 1)

        ps = StageSolver(plad, pitems, pcost, full=False, seed=seed)
        passign, _ = ps.solve(restarts=max(2, restarts // 2), iters=iters // 2)
        placed = [Placed(BACK, "back", 0, plad.points[0])] + [
            Placed(w, kind, s, plad.points[passign[w]]) for (kind, w), s in zip(words, stages)
        ]
        pages[name] = CorePage(name, spec["label"], spec["role"], psizes, placed, {})
    lay = CoreLayout(sizes, core, pages)
    _fill_at(lay)
    return dump_core(lay, f"seed {seed}, {restarts} restarts, {iters} iterations", total)


def _fill_at(lay: CoreLayout) -> None:
    for p in lay.core:
        p.at = {s.name: place(p.point, s) for s in lay.sizes[p.stage:]}
    for page in lay.pages.values():
        for p in page.items:
            p.at = {s.name: place(p.point, s) for s in page.sizes[p.stage:]}


def _grid_text(size: Size, cells: dict[tuple[int, int], str], width: int = 12) -> str:
    rows = []
    for r in range(size.rows):
        row = []
        for c in range(size.cols):
            t = cells.get((r, c), ".")
            row.append((t[: width - 1] + "~") if len(t) > width else t.ljust(width))
        rows.append("| " + " | ".join(row).rstrip() + " |")
    return "\n".join(rows)


def core_grid(lay: CoreLayout, size: Size, page: str | None = None, show_stage_upto: int | None = None) -> str:
    cells = {}
    items = lay.core if page is None else lay.pages[page].items
    for p in items:
        if size.name in p.at:
            label = p.key + ("▸" if p.kind == "page" else "")
            cells[p.at[size.name]] = label
    return _grid_text(size, cells)


def dump_core(lay: CoreLayout, how: str, total: float) -> str:
    lines = [
        "# Growing Core — layout. Generated by tools/pagesets/solve.py core "
        f"({how}; region cost {total:.4f}), then reviewed.",
        "# point: [v, u] — where the item lives (0..1 down, 0..1 across). at: its cell [row, col]",
        "# at each size from its stage on: always the cell that contains the point (pslib/grid.py).",
        "# build.py checks both against the design and the rules; edit a point, and every `at` with it.",
        "",
        "core:",
    ]
    for p in lay.core:
        lines.append(f"  {_q(p.key)}: {{stage: {lay.sizes[p.stage].name}, point: {fmt_point(p.point)}, at: {_at(p.at)}}}")
    lines.append("")
    lines.append("pages:")
    for name, page in lay.pages.items():
        lines.append(f"  {_q(name)}:")
        for p in page.items:
            lines.append(
                f"    {_q(p.key)}: {{stage: {page.sizes[p.stage].name}, point: {fmt_point(p.point)}, at: {_at(p.at)}}}"
            )
    lines.append("")
    lines.append("# For review: the core at every size (▸ opens a page).")
    lines.append("grids:")
    for s in lay.sizes:
        lines.append(f"  {s.name}: |")
        for row in core_grid(lay, s).splitlines():
            lines.append("    " + row)
    return "\n".join(lines) + "\n"


def _q(key: str) -> str:
    return json.dumps(key, ensure_ascii=False) if _needs_quote(key) else key


def _needs_quote(key: str) -> bool:
    return (
        key.lower() in {"yes", "no", "on", "off", "true", "false", "y", "n", "null", "~"}
        or any(ch in key for ch in ":'\"#{}[],&*?|<>=!%@`+")
        or key != key.strip()
    )


def _at(at: dict[str, tuple[int, int]]) -> str:
    return "{" + ", ".join(f"{k}: [{r}, {c}]" for k, (r, c) in at.items()) + "}"


def load_core(data: Data) -> CoreLayout:
    """The committed layout, checked against the design and the rules."""
    d, sizes, core_items = _core_design(data)
    raw = load_yaml(HERE / "data" / "growing-core.layout.yaml")
    names = [s.name for s in sizes]
    problems = []
    core = []
    if set(raw["core"]) != {k for k, _, _ in core_items}:
        problems.append("the layout's core items differ from the design's")
    for key, kind, k in core_items:
        r = raw["core"].get(key)
        if r is None:
            continue
        if r["stage"] != names[k]:
            problems.append(f"core {key}: stage {r['stage']} in the layout, {names[k]} in the design")
        core.append(Placed(key, kind, k, tuple(r["point"]), {s: tuple(v) for s, v in r["at"].items()}))
    opener_stage = {key: k for key, kind, k in core_items if kind == "page"}
    pages = {}
    for name, spec in d["pages"].items():
        psizes = sizes[opener_stage[name]:]
        pnames = [s.name for s in psizes]
        words = [("modifier", m["label"]) for m in spec.get("modifiers", [])] + [("word", w) for w in spec["words"]]
        stages = _page_stages(psizes, len(words))
        rp = raw["pages"].get(name, {})
        expect = [BACK] + [w for _, w in words]
        if list(rp) != expect:
            problems.append(f"page {name}: the layout's items differ from the design's (or their order)")
        placed = []
        for (kind, w), s in zip([("back", BACK)] + words, [0] + stages):
            r = rp.get(w)
            if r is None:
                continue
            if r["stage"] != pnames[s]:
                problems.append(f"page {name}: {w}: stage {r['stage']} in the layout, {pnames[s]} by priority")
            placed.append(Placed(w, kind, s, tuple(r["point"]), {k: tuple(v) for k, v in r["at"].items()}))
        mods = {m["label"]: m["modifier"] for m in spec.get("modifiers", [])}
        pages[name] = CorePage(name, spec["label"], spec["role"], psizes, placed, mods)
    lay = CoreLayout(sizes, core, pages)
    problems += check_core(lay, core_anchors(d))
    if problems:
        raise SystemExit("growing-core.layout.yaml:\n  " + "\n  ".join(problems))
    return lay


def check_core(lay: CoreLayout, lay_anchors: dict) -> list[str]:
    problems = []

    def check(where: str, items: list[Placed], sizes: list[Size], full: bool):
        for p in items:
            want = {s.name: place(p.point, s) for s in sizes[p.stage:]}
            if p.at != want:
                problems.append(f"{where} {p.key}: `at` is not where its point falls: {p.at} vs {want}")
        problems.extend(
            f"{where}: {x}" for x in check_stages(sizes, {p.key: (p.point, p.stage) for p in items}, full)
        )

    check("core", lay.core, lay.sizes, True)
    for p in lay.core:
        for size, cell in lay_anchors.get(p.key, []):
            if p.at.get(size.name) != cell:
                problems.append(f"core {p.key}: not where the designed {size} board has it ({cell})")
    for name, page in lay.pages.items():
        check(f"page {name}", page.items, page.sizes, False)
        back = next(p for p in page.items if p.kind == "back")
        if any(cell != (0, 0) for cell in back.at.values()) or back.stage != 0:
            problems.append(f"page {name}: Back is not in the top-left corner at every size")
    return problems


# ════════════════════════════════════════════════════════════════════════
# Fixed Paths
# ════════════════════════════════════════════════════════════════════════


@dataclass
class MotorSize:
    label: str  # "30"
    size: Size
    home: dict[str, tuple[int, int]]  # word or page key -> cell
    pages: dict[str, dict]  # page -> {close, words: {w: cell}, groups: {g: cell}}
    groups: dict[str, dict]  # group -> {page, close, words: {w: cell}}


@dataclass
class MotorLayout:
    sizes: list[MotorSize]
    points: dict[str, tuple[tuple[float, float], int]]  # home item -> (point, stage)
    weights: dict[str, float]


def motor_design(data: Data, aac_boost: float | None = None):
    d = data.fixed_paths
    boost = d["aac_boost"] if aac_boost is None else aac_boost
    labels = list(d["sizes"])
    sizes = [Size.parse(d["sizes"][k]) for k in labels]
    page_of, group_of = {}, {}
    for page, spec in d["pages"].items():
        for w in spec.get("words", []):
            if w in page_of:
                raise SystemExit(f"fixed-paths: {w!r} is on two pages ({page_of[w]}, {page}): one path per word")
            page_of[w] = page
        for g, ws in spec.get("groups", {}).items():
            for w in ws:
                if w in page_of:
                    raise SystemExit(f"fixed-paths: {w!r} is on two pages ({page_of[w]}, {page})")
                page_of[w] = page
                group_of[w] = g
    weights = {w: data.weight(data.entry(w), boost, d["phrase_share"]) for w in page_of}
    return d, labels, sizes, page_of, group_of, weights


def _home_words(d, sizes, weights) -> list[list[str]]:
    """The words on Home at each size, after the page keys: at the smallest,
    the beginners' words when the design names them, else the heaviest; at
    each larger size, those and then the heaviest of the rest."""
    n_pages = len(d["pages"])
    ranked = sorted(weights, key=lambda w: (-weights[w], w))
    first = list(d.get("beginner_home") or [])
    room = sizes[0].cells - n_pages
    if first:
        unknown = [w for w in first if w not in weights]
        if unknown or len(first) != room or len(set(first)) != len(first):
            raise SystemExit(f"fixed-paths beginner_home: {room} distinct words of the set needed; unknown: {unknown}")
    homes = [first or ranked[:room]]
    for s in sizes[1:]:
        have = list(homes[-1])
        have += [w for w in ranked if w not in have][: s.cells - n_pages - len(have)]
        homes.append(have)
    return homes


def _place_page(members: list[str], weights: dict[str, float], blend: dict, close: tuple[int, int], size: Size, group_keys: dict[str, float]):
    """Members and group keys on the cheapest cells, heaviest first."""
    cells = sorted((c for c in blend if c != close), key=lambda c: (blend[c], c))
    things = sorted([(w, weights[w], "word") for w in members] + [(g, wt, "group") for g, wt in group_keys.items()], key=lambda x: (-x[1], x[0]))
    if len(things) > len(cells):
        raise SystemExit(f"a page at {size} has {len(things)} items for {len(cells)} cells")
    words, groups = {}, {}
    for (name, _, kind), cell in zip(things, cells):
        (words if kind == "word" else groups)[name] = cell
    return words, groups


def _pages_at(d, size: Size, home: dict[str, tuple[int, int]], page_of, group_of, weights, blend):
    pages, groups = {}, {}
    cap = size.cells - 1
    for page, spec in d["pages"].items():
        close = home[page]
        keep = set(spec.get("always_direct", []))
        members = sorted(
            [w for w, p in page_of.items() if p == page and w not in home], key=lambda w: (w not in keep, -weights[w], w)
        )
        mods = [m["label"] for m in spec.get("modifiers", [])]
        room = cap - len(mods)
        # A full page frees places one at a time, at the least weight: opening
        # a group moves its two lightest words behind its key (two words out,
        # one key in), and an open group takes its next lightest word. So the
        # heaviest words stay on the page, the lightest of a meaning go behind
        # its key, and a key never stands for a single word.
        pool: dict[str, list[str]] = {}
        for w in members:
            if w not in keep and group_of.get(w):
                pool.setdefault(group_of[w], []).append(w)
        for ws in pool.values():
            ws.sort(key=lambda w: (weights[w], w))  # lightest first
        taken: dict[str, int] = {}
        need = len(members) - room
        while need > 0:
            best = None
            for g, ws in pool.items():
                k = taken.get(g, 0)
                if k == 0 and len(ws) >= 2:
                    step = (weights[ws[0]] + weights[ws[1]], g, 2)
                elif 0 < k < len(ws):
                    step = (weights[ws[k]], g, 1)
                else:
                    continue
                if best is None or step[:2] < best[:2]:
                    best = step
            if best is None:
                raise SystemExit(f"fixed-paths {size}: {page} cannot fit, even with its groups behind keys")
            taken[best[1]] = taken.get(best[1], 0) + best[2]
            need -= 1
        overflow = [w for g, k in taken.items() for w in pool[g][:k]]
        direct = [w for w in members if w not in overflow]
        gkeys = {}
        for w in overflow:
            gkeys[group_of[w]] = gkeys.get(group_of[w], 0.0) + weights[w]
        wts = dict(weights)
        for m in mods:
            wts[m] = 0.0
        words, gcells = _place_page(direct + mods, wts, blend, close, size, gkeys)
        pages[page] = {"close": close, "words": words, "groups": gcells}
        for g, gcell in gcells.items():
            gm = [w for w in overflow if group_of[w] == g]
            gw, _ = _place_page(gm, weights, blend, gcell, size, {})
            groups[g] = {"page": page, "close": gcell, "words": gw}
    return pages, groups


def solve_motor(data: Data, seed: int, restarts: int, iters: int) -> str:
    d, labels, sizes, page_of, group_of, weights = motor_design(data)
    ladder = ladder_for(tuple(s.name for s in sizes))
    homes = _home_words(d, sizes, weights)
    blends = [costs.blended(s, d["blend"][lab]) for lab, s in zip(labels, sizes)]

    def page_weight(page: str, k: int) -> float:
        return sum(w for x, w in weights.items() if page_of[x] == page and x not in homes[k])

    items = [Item(p, 0, data={"page": True}) for p in d["pages"]]
    seen: set[str] = set()
    for k, ws in enumerate(homes):
        for w in ws:
            if w not in seen:
                items.append(Item(w, k))
                seen.add(w)

    def cost(it: Item, atom: int) -> float:
        total = 0.0
        for k in range(it.stage, len(sizes)):
            cell = place(ladder.points[atom], sizes[k])
            wt = page_weight(it.key, k) if it.data.get("page") else weights[it.key]
            total += wt * blends[k][cell]
        return total

    solver = StageSolver(ladder, items, cost, full=True, seed=seed)
    assign, total = solver.solve(restarts=restarts, iters=iters)
    points = {it.key: (ladder.points[assign[it.key]], it.stage) for it in items}
    out_sizes = []
    for k, (lab, s) in enumerate(zip(labels, sizes)):
        home = {key: place(pt, s) for key, (pt, st) in points.items() if st <= k}
        pages, groups = _pages_at(d, s, home, page_of, group_of, weights, blends[k])
        out_sizes.append(MotorSize(lab, s, home, pages, groups))
    lay = MotorLayout(out_sizes, points, weights)
    return dump_motor(lay, labels, list(d["pages"]), f"seed {seed}, {restarts} restarts, {iters} iterations", total)


def dump_motor(lay: MotorLayout, labels, page_names, how: str, total) -> str:
    lines = [
        "# Fixed Paths — layout. Generated by tools/pagesets/solve.py motor "
        f"({how}; weighted motor cost {total:.6f}), then reviewed.",
        "# home: where each Home item lives ([v, u] point, first size). Its cell at every size is",
        "# the one containing the point. sizes: every page, as solved: on a page, `close` is the",
        "# place of the key that opened it (Back, to the board the key is on), and the heaviest words take the",
        "# cheapest cells. build.py recomputes the pages from the design and checks they match.",
        "",
        "home:",
    ]
    for key, (pt, st) in lay.points.items():
        lines.append(f"  {_q(key)}: {{from: \"{labels[st]}\", point: {fmt_point(pt)}}}")
    lines.append("")
    lines.append("sizes:")
    for ms in lay.sizes:
        lines.append(f"  \"{ms.label}\":")
        lines.append(f"    grid: {ms.size.name}")
        lines.append("    home: " + _cells(ms.home))
        lines.append("    pages:")
        for page, spec in ms.pages.items():
            lines.append(f"      {_q(page)}:")
            lines.append(f"        close: [{spec['close'][0]}, {spec['close'][1]}]")
            lines.append("        words: " + _cells(spec["words"]))
            if spec["groups"]:
                lines.append("        groups: " + _cells(spec["groups"]))
        if ms.groups:
            lines.append("    groups:")
            for g, spec in ms.groups.items():
                lines.append(f"      {_q(g)}:")
                lines.append(f"        page: {_q(spec['page'])}")
                lines.append(f"        close: [{spec['close'][0]}, {spec['close'][1]}]")
                lines.append("        words: " + _cells(spec["words"]))
    lines.append("")
    lines.append("# For review: Home at every size (▸ opens a page).")
    lines.append("grids:")
    for ms in lay.sizes:
        lines.append(f"  \"{ms.label}\": |")
        cells = {cell: key + ("▸" if key in page_names else "") for key, cell in ms.home.items()}
        for row in _grid_text(ms.size, cells, 11).splitlines():
            lines.append("    " + row)
    return "\n".join(lines) + "\n"


def _cells(m: dict[str, tuple[int, int]]) -> str:
    return "{" + ", ".join(f"{_q(k)}: [{r}, {c}]" for k, (r, c) in m.items()) + "}"


def load_motor(data: Data) -> MotorLayout:
    """The committed layout, recomputed from the design and compared."""
    d, labels, sizes, page_of, group_of, weights = motor_design(data)
    raw = load_yaml(HERE / "data" / "fixed-paths.layout.yaml")
    homes = _home_words(d, sizes, weights)
    problems = []
    points = {}
    for key, r in raw["home"].items():
        points[key] = (tuple(r["point"]), labels.index(r["from"]))
    expect_keys = set(d["pages"]) | {w for ws in homes for w in ws}
    if set(points) != expect_keys:
        problems.append(f"home items differ from the design: {sorted(set(points) ^ expect_keys)}")
    for k, ws in enumerate(homes):
        for w in ws:
            first = next(i for i, h in enumerate(homes) if w in h)
            if w in points and points[w][1] != first:
                problems.append(f"{w}: on Home from {labels[points[w][1]]} in the layout, {labels[first]} by weight")
    problems += check_stages(sizes, points, full=True)
    out_sizes = []
    for k, (lab, s) in enumerate(zip(labels, sizes)):
        home = {key: place(pt, s) for key, (pt, st) in points.items() if st <= k}
        rs = raw["sizes"][lab]
        if {kk: tuple(v) for kk, v in rs["home"].items()} != home:
            problems.append(f"{lab}: Home cells in the layout are not where the points fall")
        pages, groups = _pages_at(d, s, home, page_of, group_of, weights, costs.blended(s, d["blend"][lab]))
        for page, spec in pages.items():
            rp = rs["pages"][page]
            if tuple(rp["close"]) != spec["close"] or {w: tuple(c) for w, c in rp["words"].items()} != spec["words"] or {g: tuple(c) for g, c in rp.get("groups", {}).items()} != spec["groups"]:
                problems.append(f"{lab}: page {page} differs from the design's placement")
        for g, spec in groups.items():
            rg = rs.get("groups", {}).get(g)
            if rg is None or {w: tuple(c) for w, c in rg["words"].items()} != spec["words"]:
                problems.append(f"{lab}: group {g} differs from the design's placement")
        out_sizes.append(MotorSize(lab, s, home, pages, groups))
    if problems:
        raise SystemExit("fixed-paths.layout.yaml:\n  " + "\n  ".join(problems))
    return MotorLayout(out_sizes, points, weights)