render.py

"""Poster rendering: layer composition, typography, saving."""

import os
from dataclasses import dataclass
from typing import IO, Union

import matplotlib
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np
import osmnx as ox
from geopandas import GeoDataFrame
from matplotlib.figure import Figure
from matplotlib.font_manager import FontProperties
from networkx import MultiDiGraph
from shapely.geometry import Point

from layout import (
    BASE_ATTR,
    BASE_COORDS,
    BASE_MAIN,
    BASE_SUB,
    LogoSpec,
    compute_text_layout,
    fit_title_size,
    format_coords,
    is_latin_script,
    place_logo,
)
from osm_fetch import fetch_all
from presets import SizeSpec
from progress import ProgressFn, emit

# Rendering is headless and never opens a window. Agg also guarantees
# fig.canvas.get_renderer(), which fit_title_size needs to measure the title.
matplotlib.use("Agg")

PREVIEW_DPI = 96

ATTRIBUTION_TEXT = "© OpenStreetMap contributors"

# Text metrics, the divider rule and now road strokes all scale against a 12 inch
# reference side. Road widths are matplotlib points (1/72 in) -- physical units --
# so without this a 30x40 in poster draws the same 0.4 pt residential street as a
# 12x16 in one and the whole network visually disappears.
REFERENCE_SIDE_IN = 12.0


def poster_scale_factor(width_in, height_in):
    """Stroke and type scale for a poster, relative to the 12 in reference side."""
    return min(width_in, height_in) / REFERENCE_SIDE_IN


def compensated_distance(dist, width_in, height_in) -> int:
    """
    Fetch/crop radius in whole metres for a poster of this aspect ratio.

    get_crop_limits cuts the fetched box inward to the figure aspect, so the
    fetch is widened by the aspect ratio first. The result lands verbatim in the
    OSM cache key (osm_fetch.fetch_graph / fetch_features), which is why it is
    rounded: --pixels 3000x4000 computes 666.6666666666667 where the cached
    12x16 entry holds 666.6666666666666, and the identical poster refetches.

    The /4 is long-standing semantics -- `-d 2000` renders a ~666 m radius --
    and is deliberately preserved.
    """
    return round(dist * (max(height_in, width_in) / min(height_in, width_in)) / 4)


def create_gradient_fade(ax, color, location="bottom", extent_frac=0.25, zorder=10,
                         opacity=1.0):
    """
    Fade the map into a flat color at the top or bottom edge of the axes.

    Args:
        ax: Axes to draw on
        color: Fade color (the theme's gradient_color)
        location: "bottom" or "top" edge
        extent_frac: Fraction of the axes height the fade covers
        zorder: Draw order; must sit above the map and below the text
        opacity: Scales the ramp's maximum alpha, 0..1. The shape of the ramp is
            unchanged -- it still runs linearly from the flat edge to nothing --
            so a half-opacity fade is the same gradient seen through the map
            rather than a differently-shaped one. 1.0 is the legacy fade.
    """
    vals = np.linspace(0, 1, 256).reshape(-1, 1)
    gradient = np.hstack((vals, vals))

    rgb = mcolors.to_rgb(color)
    my_colors = np.zeros((256, 4))
    my_colors[:, 0] = rgb[0]
    my_colors[:, 1] = rgb[1]
    my_colors[:, 2] = rgb[2]

    if location == "bottom":
        my_colors[:, 3] = np.linspace(opacity, 0, 256)
        extent_y_start = 0.0
        extent_y_end = extent_frac
    else:
        my_colors[:, 3] = np.linspace(0, opacity, 256)
        extent_y_start = 1.0 - extent_frac
        extent_y_end = 1.0

    custom_cmap = mcolors.ListedColormap(my_colors)

    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    y_range = ylim[1] - ylim[0]

    y_bottom = ylim[0] + y_range * extent_y_start
    y_top = ylim[0] + y_range * extent_y_end

    ax.imshow(
        gradient,
        extent=[xlim[0], xlim[1], y_bottom, y_top],
        aspect="auto",
        cmap=custom_cmap,
        zorder=zorder,
        origin="lower",
    )


def get_edge_colors_by_type(g, theme):
    """
    Assigns colors to edges based on road type hierarchy.
    Returns a list of colors corresponding to each edge in the graph.
    """
    edge_colors = []

    for _u, _v, data in g.edges(data=True):
        # Get the highway type (can be a list or string)
        highway = data.get("highway", "unclassified")

        # Handle list of highway types (take the first one)
        if isinstance(highway, list):
            highway = highway[0] if highway else "unclassified"

        # Assign color based on road type
        if highway in ["motorway", "motorway_link"]:
            color = theme["road_motorway"]
        elif highway in ["trunk", "trunk_link", "primary", "primary_link"]:
            color = theme["road_primary"]
        elif highway in ["secondary", "secondary_link"]:
            color = theme["road_secondary"]
        elif highway in ["tertiary", "tertiary_link"]:
            color = theme["road_tertiary"]
        elif highway in ["residential", "living_street", "unclassified"]:
            color = theme["road_residential"]
        else:
            color = theme["road_default"]

        edge_colors.append(color)

    return edge_colors


def get_edge_widths_by_type(g, scale=1.0):
    """
    Assigns line widths to edges based on road type.
    Major roads get thicker lines.

    `scale` keeps road weight constant *relative to the poster* -- pass
    poster_scale_factor(width_in, height_in). See REFERENCE_SIDE_IN.
    """
    edge_widths = []

    for _u, _v, data in g.edges(data=True):
        highway = data.get("highway", "unclassified")

        if isinstance(highway, list):
            highway = highway[0] if highway else "unclassified"

        # Assign width based on road importance
        if highway in ["motorway", "motorway_link"]:
            width = 1.2
        elif highway in ["trunk", "trunk_link", "primary", "primary_link"]:
            width = 1.0
        elif highway in ["secondary", "secondary_link"]:
            width = 0.8
        elif highway in ["tertiary", "tertiary_link"]:
            width = 0.6
        else:
            width = 0.4

        edge_widths.append(width * scale)

    return edge_widths


def get_crop_limits(g_proj, center_lat_lon, fig, dist):
    """
    Crop inward to preserve aspect ratio while guaranteeing
    full coverage of the requested radius.
    """
    lat, lon = center_lat_lon

    # Project center point into graph CRS
    center = (
        ox.projection.project_geometry(
            Point(lon, lat),
            crs="EPSG:4326",
            to_crs=g_proj.graph["crs"]
        )[0]
    )
    center_x, center_y = center.x, center.y

    fig_width, fig_height = fig.get_size_inches()
    aspect = fig_width / fig_height

    # Start from the *requested* radius
    half_x = dist
    half_y = dist

    # Cut inward to match aspect
    if aspect > 1:  # landscape → reduce height
        half_y = half_x / aspect
    else:  # portrait → reduce width
        half_x = half_y * aspect

    return (
        (center_x - half_x, center_x + half_x),
        (center_y - half_y, center_y + half_y),
    )


def gradient_locations(lay, fade_enabled=True):
    """
    Edges that get a fade for a given text layout.

    Both edges whenever a text block is visible (top or bottom) — that framing is
    the legacy look and the poster reads as unbalanced without it. The wallpaper
    case (text_position="none") gets no fade at all: with no text to sit on, the
    scrim has nothing to do.

    `fade_enabled=False` is UX spec §2.4's Edge fade toggle turned off, and it
    answers the same way: no edges, so no gradient artists are created at all
    rather than transparent ones. The text block still renders — that is the
    whole difference between switching the fade off and wallpaper mode.
    """
    return () if lay is None or not fade_enabled else ("bottom", "top")


def _font(font_path, size, fallback_weight="normal"):
    """FontProperties for a bundled/downloaded font, or a monospace fallback."""
    if font_path:
        return FontProperties(fname=font_path, size=size)
    return FontProperties(family="monospace", weight=fallback_weight, size=size)


def _plot_polygons(gdf, ax, g_proj, facecolor, zorder):
    """Project and draw the polygon rows of a feature GeoDataFrame."""
    if gdf is None or gdf.empty:
        return
    # Filter to only polygon/multipolygon geometries to avoid point features showing as dots
    polys = gdf[gdf.geometry.type.isin(["Polygon", "MultiPolygon"])]
    if polys.empty:
        return
    # Project features into the same CRS as the graph
    try:
        polys = ox.projection.project_gdf(polys)
    except Exception:
        polys = polys.to_crs(g_proj.graph["crs"])
    polys.plot(ax=ax, facecolor=facecolor, edgecolor="none", zorder=zorder)


@dataclass(frozen=True)
class FetchedData:
    """Everything one OSM job produces, ready to render any number of times."""

    graph: MultiDiGraph | None
    water: GeoDataFrame | None
    parks: GeoDataFrame | None
    point: tuple[float, float]
    graph_proj: MultiDiGraph | None = None


@dataclass(frozen=True)
class RenderOptions:
    """Everything that turns FetchedData into one specific poster."""

    theme: dict
    size: SizeSpec
    dist: int
    display_city: str
    display_country: str
    text_position: str = "bottom"
    logo: LogoSpec | None = None
    fonts: dict | None = None
    # UX spec §2.4's Edge fade. The three defaults are the legacy fade exactly:
    # both edges, the 4 in extent, full alpha. Nothing here applies in wallpaper
    # mode -- see gradient_locations.
    fade_enabled: bool = True
    fade_height: float = 1.0
    fade_opacity: float = 1.0


def _empty_map_frame(ax, width_in, height_in, dist):
    """
    Axes for the placeholder poster: real aspect, real crop maths, no roads.

    Mirrors get_crop_limits' inward cut so the empty state has the same framing
    a loaded city would, centred on (0, 0). No fake roads are drawn -- inventing
    a city is a lie the user notices the moment they load a real one.
    """
    half_x = half_y = float(dist)
    aspect = width_in / height_in
    if aspect > 1:
        half_y = half_x / aspect
    else:
        half_x = half_y * aspect
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlim(-half_x, half_x)
    ax.set_ylim(-half_y, half_y)
    ax.set_axis_off()


def render_poster_figure(
    data: FetchedData, options: RenderOptions, progress: ProgressFn | None = None
) -> Figure:
    """
    Assemble one poster figure from already-fetched data.

    Pure with respect to the network and the filesystem: no fetching, no saving.
    The caller owns the returned figure and must plt.close() it.

    data.graph may be None -- that renders the placeholder poster (real
    typography and theme colours over an empty map frame).
    data.graph_proj, when supplied, skips the expensive ox.project_graph call,
    which is what lets the server re-render on a theme switch in under a second.
    """
    theme = options.theme
    width, height = options.size.width_in, options.size.height_in
    scale_factor = poster_scale_factor(width, height)
    crop_dist = compensated_distance(options.dist, width, height)

    fig, ax = plt.subplots(figsize=(width, height), facecolor=theme["bg"])
    ax.set_facecolor(theme["bg"])
    # fit_title_size measures the title against the figure width, which is only
    # equivalent to the axes width while the axes fills the figure.
    ax.set_position((0.0, 0.0, 1.0, 1.0))

    if data.graph is None:
        emit(progress, "render_start", edge_count=0)
        _empty_map_frame(ax, width, height, crop_dist)
    else:
        # Project to a metric CRS so distances and aspect are linear (metres).
        g_proj = data.graph_proj if data.graph_proj is not None else ox.project_graph(data.graph)
        emit(progress, "render_start", edge_count=g_proj.number_of_edges())

        # Layer 1: water and park polygons
        _plot_polygons(data.water, ax, g_proj, theme["water"], zorder=0.5)
        _plot_polygons(data.parks, ax, g_proj, theme["parks"], zorder=0.8)

        # Layer 2: roads with hierarchy coloring
        edge_colors = get_edge_colors_by_type(g_proj, theme)
        edge_widths = get_edge_widths_by_type(g_proj, scale=scale_factor)

        crop_xlim, crop_ylim = get_crop_limits(g_proj, data.point, fig, crop_dist)
        ox.plot_graph(
            g_proj, ax=ax, bgcolor=theme["bg"], node_size=0,
            edge_color=edge_colors, edge_linewidth=edge_widths,
            show=False, close=False,
        )
        ax.set_aspect("equal", adjustable="box")
        ax.set_xlim(crop_xlim)
        ax.set_ylim(crop_ylim)

    lay = compute_text_layout(width, height, options.text_position,
                              fade_height=options.fade_height)
    fonts = options.fonts

    # Layer 3: gradients (see gradient_locations for which edges and why)
    for edge in gradient_locations(lay, options.fade_enabled):
        create_gradient_fade(ax, theme["gradient_color"], location=edge,
                             extent_frac=lay.gradient_frac, zorder=10,
                             opacity=options.fade_opacity)

    if lay is not None:
        font_path_bold = fonts["bold"] if fonts else None
        font_path_light = fonts["light"] if fonts else None
        font_path_regular = fonts["regular"] if fonts else None

        # Latin scripts get uppercase letter-spacing ("P  A  R  I  S"); other
        # scripts (CJK, Thai, Arabic, ...) are left exactly as supplied.
        if is_latin_script(options.display_city):
            title_text = "  ".join(list(options.display_city.upper()))
        else:
            title_text = options.display_city

        title_size = fit_title_size(fig, ax, title_text, font_path_bold,
                                    BASE_MAIN * scale_factor)

        ax.text(
            0.5, lay.title_y, title_text, transform=ax.transAxes,
            color=theme["text"], ha="center",
            fontproperties=_font(font_path_bold, title_size, "bold"), zorder=11,
        )
        ax.plot(
            [0.4, 0.6], [lay.divider_y, lay.divider_y], transform=ax.transAxes,
            color=theme["text"], linewidth=1 * scale_factor, zorder=11,
        )
        ax.text(
            0.5, lay.country_y, options.display_country.upper(), transform=ax.transAxes,
            color=theme["text"], ha="center",
            fontproperties=_font(font_path_light, BASE_SUB * scale_factor), zorder=11,
        )
        ax.text(
            0.5, lay.coords_y, format_coords(*data.point), transform=ax.transAxes,
            color=theme["text"], alpha=0.7, ha="center",
            fontproperties=_font(font_path_regular, BASE_COORDS * scale_factor), zorder=11,
        )

    # Attribution is always drawn: OpenStreetMap data is ODbL, credit is required.
    ax.text(
        0.98, 0.02, ATTRIBUTION_TEXT, transform=ax.transAxes,
        color=theme["text"], alpha=0.5, ha="right", va="bottom",
        fontproperties=_font(fonts["light"] if fonts else None, BASE_ATTR * scale_factor),
        zorder=11,
    )

    if options.logo is not None:
        # above-title anchors to the bottom text block; with top or no text,
        # fall back to the bottom-area default (title_y=None -> 0.07 base).
        anchor_y = lay.title_y if (lay is not None and not lay.at_top) else None
        place_logo(fig, options.logo, width, height, title_y=anchor_y)

    return fig


def save_figure(
    fig: Figure,
    target: Union[str, "os.PathLike", "IO[bytes]"],
    size: SizeSpec,
    fmt: str,
    *,
    facecolor: str,
    preview: bool = False,
    progress: ProgressFn | None = None,
) -> None:
    """
    Write a figure to a path or a binary buffer.

    `target` is anything matplotlib's savefig accepts: a str or Path for the
    export path, an io.BytesIO for the preview path. The figure is left open --
    the caller closes it.

    `fmt` is case-insensitive; a caller passing "PNG" still gets the dpi branch.
    """
    fmt = fmt.lower()
    emit(progress, "save_start", fmt=fmt, preview=preview)
    # Deliberately untyped values: dpi is an int and pad_inches a float, so a
    # dict[str, str] inferred from facecolor alone would not accept them.
    save_kwargs: dict = {"facecolor": facecolor}
    if not size.exact_pixels:
        save_kwargs.update(bbox_inches="tight", pad_inches=0.05)
    if fmt == "png":
        save_kwargs["dpi"] = PREVIEW_DPI if preview else size.dpi
    fig.savefig(target, format=fmt, **save_kwargs)


def create_poster(
    city,
    country,
    point,
    dist,
    output_file,
    output_format,
    theme,
    size: SizeSpec,
    *,
    network_type="all",
    text_position="bottom",
    logo: LogoSpec | None = None,
    preview=False,
    display_city=None,
    display_country=None,
    country_label=None,
    fonts=None,
    fade_enabled=True,
    fade_height=1.0,
    fade_opacity=1.0,
) -> None:
    """
    Generate a complete map poster with roads, water, parks, and typography.

    Thin wrapper over fetch_all -> render_poster_figure -> save_figure, kept for
    the CLI: same arguments, same console output, same files on disk.

    Args:
        city: City name used for geocoding context and console output
        country: Country name used for console output
        point: (latitude, longitude) tuple for map center
        dist: Map radius in meters
        output_file: Path where the poster will be saved
        output_format: File format ('png', 'svg', or 'pdf')
        theme: Validated theme dict (see themes.load_theme)
        size: Resolved SizeSpec (inches, dpi, exact-pixel flag)
        network_type: OSMnx network type for the street graph
        text_position: 'bottom', 'top', or 'none'
        logo: Optional LogoSpec to overlay
        preview: Render PNGs at PREVIEW_DPI instead of the full size dpi
        display_city: Optional override for the city text on the poster
        display_country: Optional override for the country text on the poster
        country_label: Legacy override for the country text (lower priority)
        fonts: Dict with 'light', 'regular', 'bold' font paths, or None
        fade_enabled: Draw the edge gradients at all (--no-fade turns them off)
        fade_height: Multiplier on the standard 4 in fade extent, 0.25-2.0
        fade_opacity: Scale on the fade's maximum alpha, 0.0-1.0

    Raises:
        RuntimeError: If street network data cannot be retrieved
    """
    display_city = display_city or city
    display_country = display_country or country_label or country

    print(f"\nGenerating map for {city}, {country}...")
    width, height = size.width_in, size.height_in
    # Compensate for the viewport crop applied by get_crop_limits
    compensated_dist = compensated_distance(dist, width, height)

    g, water, parks = fetch_all(point, compensated_dist, network_type=network_type)
    if g is None:
        raise RuntimeError("Failed to retrieve street network data.")
    print("✓ All data retrieved successfully!")

    print("Rendering map...")
    print("Applying road hierarchy colors...")
    fig = render_poster_figure(
        FetchedData(graph=g, water=water, parks=parks, point=tuple(point)),
        RenderOptions(
            theme=theme, size=size, dist=dist,
            display_city=display_city, display_country=display_country,
            text_position=text_position, logo=logo, fonts=fonts,
            fade_enabled=fade_enabled, fade_height=fade_height,
            fade_opacity=fade_opacity,
        ),
    )
    try:
        print(f"Saving to {output_file}...")
        # save_figure lower-cases fmt itself, so the CLI hands it over verbatim.
        save_figure(fig, output_file, size, output_format,
                    facecolor=theme["bg"], preview=preview)
    finally:
        plt.close(fig)
    print(f"✓ Done! Poster saved as {output_file}")