#!/usr/bin/env python3
"""The three pictures for the journal entry on how a citizen finds a path.

Run from the repository root:

    python research/hesper/pathfinding/figures.py

It writes into `site/public/images/`:

  hesper-path-live-2026-09-09.png   the human's own 5-versus-3 case, drawn on
                                    the live world's real terrain (fetched from
                                    the public /api/tiles, no key)
  hesper-path-lake-2026-09-09.png   a lake in the way: the search settling
                                    outward and the path it proves
  hesper-path-frontier-2026-09-09.png
                                    the same search, shaded by the order tiles
                                    were settled — what "cheapest first" looks
                                    like

Colours are Hesper's own palette (`hesper/assets/palette.json`), so a figure
here and the map on the site are the same world.

The paths in every figure come from the engine: `hesper.laws.movement.pathfind
.find_path`, the function `return_home` uses. The *frontier order* is not
something that function returns, so it is recomputed here by a local Dijkstra
with the same costs and the same four-way adjacency — and the script asserts
that the local search finds a path of the same cost as the engine's, so the
picture cannot quietly drift from the thing it illustrates.
"""

import heapq
import json
import os
import sys
import urllib.request

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import to_rgb
from matplotlib.patches import Rectangle

# .../research/hesper/pathfinding/figures.py -> the repository root is four up.
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
    os.path.dirname(os.path.abspath(__file__)))))
sys.path.insert(0, ROOT)

from hesper.laws.movement.pathfind import find_path  # noqa: E402
from hesper.laws.terrain.render import TERRAIN_COLOURS  # noqa: E402

OUT = os.path.join(ROOT, "site", "public", "images")
API = "https://hesper.untilnextsession.com/api/tiles"

# The cost table the engine charges, read from the same vocabulary file the
# engine reads rather than typed out here.
with open(os.path.join(ROOT, "hesper", "vocab", "v1.json"), encoding="utf-8") as fh:
    COSTS = next(
        m["cost_table"]
        for m in json.load(fh)["movement_modes"]
        if m["name"] == "ground"
    )

with open(os.path.join(ROOT, "hesper", "assets", "palette.json"), encoding="utf-8") as fh:
    PALETTE = {c["name"]: c["hex"] for c in json.load(fh)["colours"] if c["hex"]}

GROUND = PALETTE["dusk"]
INK = PALETTE["ivory"]
DIM = PALETTE["stone light"]
AMBER = PALETTE["star amber"]
RULE = PALETTE["slate"]

TERRAIN_HEX = {t: PALETTE[c] for t, c in TERRAIN_COLOURS.items()}


def shade(hex_colour, factor):
    """Lighten (factor > 1) or darken (factor < 1) a palette colour."""
    r, g, b = to_rgb(hex_colour)
    if factor >= 1:
        return tuple(c + (1 - c) * (factor - 1) for c in (r, g, b))
    return tuple(c * factor for c in (r, g, b))


def dijkstra_with_order(grid, start, goal):
    """The same search the engine runs, plus the order tiles were settled.

    Returns (path, cost, order) where `order` maps a tile to its settle rank.
    `path` is the list of steps after `start`, exactly as `find_path` returns.
    """
    dist = {start: 0}
    came = {}
    order = {}
    heap = [(0, start)]
    settled = set()
    while heap:
        cost, node = heapq.heappop(heap)
        if node in settled:
            continue
        settled.add(node)
        order[node] = len(order)
        if node == goal:
            path = []
            here = node
            while here != start:
                path.append([here[0], here[1]])
                here = came[here]
            path.reverse()
            return path, cost, order
        x, y = node
        for nxt in ((x + 1, y), (x - 1, y), (x, y + 1), (x, y - 1)):
            if nxt in settled or nxt not in grid:
                continue
            step = COSTS.get(grid[nxt])
            if step is None:
                continue
            through = cost + step
            if through < dist.get(nxt, float("inf")):
                dist[nxt] = through
                came[nxt] = node
                heapq.heappush(heap, (through, nxt))
    return None, None, order


def path_cost(grid, path):
    """What a path costs: every tile entered, the tile you start on free."""
    return sum(COSTS[grid[(int(x), int(y))]] for x, y in path)


def l_shape(start, goal):
    """The naive path the play page built until H59: all of x, then all of y."""
    (sx, sy), (gx, gy) = start, goal
    steps = []
    x, y = sx, sy
    while x != gx:
        x += 1 if gx > x else -1
        steps.append([x, y])
    while y != gy:
        y += 1 if gy > y else -1
        steps.append([x, y])
    return steps


def new_axes(width, height, title, subtitle):
    fig, ax = plt.subplots(figsize=(width, height))
    fig.patch.set_facecolor(GROUND)
    ax.set_facecolor(GROUND)
    ax.set_title(title, color=INK, fontsize=13, fontweight="bold",
                 fontfamily="monospace", loc="left", pad=18)
    ax.text(0, 1.012, subtitle, transform=ax.transAxes, color=DIM,
            fontsize=9, fontfamily="monospace", va="bottom")
    ax.set_xticks([])
    ax.set_yticks([])
    for spine in ax.spines.values():
        spine.set_visible(False)
    return fig, ax


def draw_grid(ax, grid, xs, ys, order=None, label_costs=True):
    """One square per tile, in the world's own colours. y grows downward."""
    for (x, y), terrain in grid.items():
        base = TERRAIN_HEX[terrain]
        if order is not None:
            # Settled tiles keep their colour; unvisited ones go flat and dark,
            # so the picture shows how little of the map the search touched.
            colour = shade(base, 1.0) if (x, y) in order else shade(base, 0.45)
        else:
            colour = base
        ax.add_patch(Rectangle((x - 0.5, -y - 0.5), 1, 1, facecolor=colour,
                               edgecolor=GROUND, linewidth=1.2))
        if label_costs and COSTS.get(terrain) is not None:
            ax.text(x, -y, str(COSTS[terrain]), ha="center", va="center",
                    color=INK, fontsize=8, fontfamily="monospace", alpha=0.55)
    ax.set_xlim(min(xs) - 0.7, max(xs) + 0.7)
    ax.set_ylim(-max(ys) - 0.7, -min(ys) + 0.7)
    ax.set_aspect("equal")


def draw_path(ax, start, path, colour, label, style="-", width=2.6, offset=0.0):
    xs = [start[0] + offset] + [p[0] + offset for p in path]
    ys = [-start[1] - offset] + [-p[1] - offset for p in path]
    ax.plot(xs, ys, style, color=colour, linewidth=width, solid_capstyle="round",
            zorder=5, label=label)
    ax.plot(xs[1:-1], ys[1:-1], "o", color=colour, markersize=5, zorder=6)


def mark(ax, xy, text, colour):
    """A ring on the tile, with its word inside the same tile.

    The word sits at the top edge of the tile it names. Putting it above the
    ring, which is what a first draft did, prints "to" inside the neighbouring
    tile and labels the wrong square.
    """
    ax.plot([xy[0]], [-xy[1]], "o", color=colour, markersize=11,
            markeredgecolor=GROUND, markeredgewidth=1.6, zorder=7)
    ax.text(xy[0], -xy[1] + 0.30, text, ha="center", va="center", color=colour,
            fontsize=8.5, fontfamily="monospace", fontweight="bold", zorder=7)


# --------------------------------------------------------------------------
# Figure 1: the live world, the human's own case
# --------------------------------------------------------------------------

def live_grid(x0, y0, w, h):
    url = f"{API}?x={x0}&y={y0}&w={w}&h={h}"
    with urllib.request.urlopen(url, timeout=20) as response:
        data = json.load(response)
    return {(t["x"], t["y"]): t["t"] for t in data["tiles"]}, data["turn"]


def figure_live():
    x0, y0, w, h = 254, 250, 7, 6
    grid, turn = live_grid(x0, y0, w, h)
    start, goal = (257, 253), (256, 252)

    naive = l_shape(start, goal)
    cheap = find_path(lambda x, y: (grid.get((x, y)), None), start, goal, COSTS)
    assert cheap is not None, "the engine found no path on the live world"

    naive_cost, cheap_cost = path_cost(grid, naive), path_cost(grid, cheap)

    fig, ax = new_axes(7.4, 6.2,
                       "The same two tiles, two ways round",
                       f"Hesper, live world at turn {turn}. The number on a tile "
                       f"is what it costs to walk onto.")
    draw_grid(ax, grid, range(x0, x0 + w), range(y0, y0 + h))
    draw_path(ax, start, naive, DIM, f"x first, then y — {naive_cost} points",
              style="--", width=2.2, offset=0.10)
    draw_path(ax, start, cheap, AMBER, f"cheapest — {cheap_cost} points",
              offset=-0.10)
    mark(ax, start, "from", INK)
    mark(ax, goal, "to", INK)

    legend = ax.legend(loc="lower center", bbox_to_anchor=(0.5, -0.09), ncol=2,
                       frameon=False, prop={"family": "monospace", "size": 9})
    for text in legend.get_texts():
        text.set_color(INK)
    ax.text(0, -0.15, "The dashed walk enters the hill (3) then the forest (2). "
                      "The amber walk enters the plain (1) then the same forest.",
            transform=ax.transAxes, color=DIM, fontsize=8.5,
            fontfamily="monospace")

    out = os.path.join(OUT, "hesper-path-live-2026-09-09.png")
    fig.savefig(out, dpi=170, bbox_inches="tight", facecolor=GROUND)
    plt.close(fig)
    return out, naive_cost, cheap_cost, turn


# --------------------------------------------------------------------------
# Figures 2 and 3: a lake in the way
# --------------------------------------------------------------------------

LAKE_MAP = [
    "ppppppppppppppp",
    "pppfffppppppppp",
    "ppffwwwwwppfffp",
    "ppfwwwwwwwpfffp",
    "pppwwwwwwwwpffp",
    "ppphwwwwwwwpppp",
    "pphhhwwwwwppppp",
    "pphhhppwwpppppp",
    "ppphhpppppppppp",
    "ppppppppppppppp",
]
LETTERS = {"p": "plain", "f": "forest", "h": "hill", "w": "water"}


def lake_grid():
    return {
        (x, y): LETTERS[c]
        for y, row in enumerate(LAKE_MAP)
        for x, c in enumerate(row)
    }


def figure_lake():
    grid = lake_grid()
    start, goal = (1, 4), (13, 4)

    path, cost, order = dijkstra_with_order(grid, start, goal)
    engine = find_path(lambda x, y: (grid.get((x, y)), None), start, goal, COSTS)
    assert engine is not None and path_cost(grid, engine) == cost, (
        "the picture's search disagrees with the engine's")

    xs = range(len(LAKE_MAP[0]))
    ys = range(len(LAKE_MAP))

    # The path, on the plain map.
    fig, ax = new_axes(9.0, 6.4,
                       "Water is a hole in the map, not a wall to climb",
                       "Four-way steps. Plain costs 1, forest 2, hill 3. Water "
                       "has no cost at all, so it is never entered.")
    draw_grid(ax, grid, xs, ys)
    draw_path(ax, start, path, AMBER, None)
    mark(ax, start, "from", INK)
    mark(ax, goal, "to", INK)
    ax.text(0, -0.07, f"The cheapest way round is {len(path)} steps and "
                      f"{cost} movement points. The straight line is 12 steps "
                      f"and would need to swim.",
            transform=ax.transAxes, color=DIM, fontsize=8.5,
            fontfamily="monospace")
    out_path = os.path.join(OUT, "hesper-path-lake-2026-09-09.png")
    fig.savefig(out_path, dpi=170, bbox_inches="tight", facecolor=GROUND)
    plt.close(fig)

    # The same search, shaded by the order tiles were settled.
    fig, ax = new_axes(9.0, 6.4,
                       "What the search actually touched",
                       "Every settled tile numbered by the order it was settled: "
                       "cheapest-so-far first, always.")
    draw_grid(ax, grid, xs, ys, order=order, label_costs=False)
    ranks = {t: i for t, i in order.items()}
    for (x, y), rank in ranks.items():
        ax.text(x, -y, str(rank), ha="center", va="center", color=INK,
                fontsize=6.5, fontfamily="monospace", alpha=0.75)
    draw_path(ax, start, path, AMBER, None, width=2.0)
    mark(ax, start, "from", INK)
    mark(ax, goal, "to", INK)
    water = sum(1 for terrain in grid.values() if COSTS.get(terrain) is None)
    ax.text(0, -0.075, f"{len(order)} tiles settled of the {len(grid) - water} "
                       f"that can be walked on; the {water} water tiles are the\n"
                       f"only ones never entered. A goal this far away costs you "
                       f"nearly everything nearer\nthan it — which is why the "
                       f"engine bounds the search.",
            transform=ax.transAxes, color=DIM, fontsize=8.5,
            fontfamily="monospace")
    out_frontier = os.path.join(OUT, "hesper-path-frontier-2026-09-09.png")
    fig.savefig(out_frontier, dpi=170, bbox_inches="tight", facecolor=GROUND)
    plt.close(fig)

    return out_path, out_frontier, path, cost, len(order), len(grid)


if __name__ == "__main__":
    live, naive_cost, cheap_cost, turn = figure_live()
    print(f"{live}: live turn {turn}, L-shape {naive_cost} points, "
          f"cheapest {cheap_cost} points")
    lake, frontier, path, cost, settled, total = figure_lake()
    print(f"{lake}: {len(path)} steps, {cost} points")
    print(f"{frontier}: {settled} tiles settled of {total}")
