#!/usr/bin/env python3
"""Render the Copenhagen transit map to PNG/SVG/PDF.

Reads data/processed/master.gpkg (layers: lines, stops, stations — all
GTFS-derived by prepare.py) + area.gpkg + config/styling.yaml, applies CLI
filters, and composes a printable map:
  - Carto/Esri grey basemap via contextily (EPSG:3857)
  - two-tone lines: dark casing + colour body, per route
  - z-order: bus (bottom) -> s_tog -> light_rail -> regional -> metro (top)
  - shapeburst fade mask outside the City Pass area
  - station labels (metro + S-tog) via adjustText

All transit data comes from the Rejseplanen GTFS feed. Stop/station layers
are pre-collapsed (one point per stop name / station) by prepare.py, so
rendering is pure plotting — no clustering or merging here.

Usage:
  uv run python render.py                       # default: all modes, PNG, markers only
  uv run python render.py --labels              # add station name labels
  uv run python render.py --no-stops            # no station markers
  uv run python render.py --no-buses --format svg --out map.svg
  uv run python render.py --bus-subset A,C,S     # only trunk buses
  uv run python render.py --modes metro,s_tog     # rail-only
"""
import argparse
from pathlib import Path

import geopandas as gpd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yaml
from affine import Affine
from adjustText import adjust_text
from scipy.ndimage import distance_transform_edt

from _common import CONFIG, OUTPUT, PROCESSED, UA

TARGET_CRS = "EPSG:3857"


def load_args():
    p = argparse.ArgumentParser(description="Render the Copenhagen transit map.")
    p.add_argument("--modes", default=None,
                   help="comma-separated styles to include (metro,s_tog,light_rail,regional,bus)")
    p.add_argument("--no-buses", action="store_true", help="omit bus layer")
    p.add_argument("--bus-subset", default=None,
                   help="comma-separated bus categories to keep (A,C,S,regular)")
    p.add_argument("--max-bus-routes", type=int, default=None,
                   help="cap on number of distinct bus refs drawn")
    p.add_argument("--dpi", type=int, default=200)
    p.add_argument("--size", default="16",
                   help="figure size in inches: 'W' (height auto) or 'WxH'")
    p.add_argument("--format", default="png", help="output format (png/svg/pdf)")
    p.add_argument("--out", default=None, help="output path (default: output/copenhagen.<fmt>)")
    p.add_argument("--no-basemap", action="store_true", help="skip basemap tiles (offline/faster)")
    p.add_argument("--labels", action="store_true", help="draw station name labels (off by default)")
    p.add_argument("--no-stops", action="store_true", help="skip station markers")
    return p.parse_args()


def parse_size(spec, aspect):
    if "x" in spec:
        w, h = spec.split("x")
        return float(w), float(h)
    w = float(spec)
    return w, w / aspect


def filter_lines(lines, args):
    keep = set(lines["style"].unique())
    if args.modes:
        keep &= {m.strip() for m in args.modes.split(",")}
    if args.no_buses:
        keep.discard("bus")
    out = lines[lines["style"].isin(keep)].copy()

    if "bus" in keep and out["style"].eq("bus").any():
        bus = out[out["style"].eq("bus")]
        other = out[~out["style"].eq("bus")]
        if args.bus_subset:
            cats = {c.strip() for c in args.bus_subset.split(",")}
            bus = bus[bus["bus_category"].isin(cats)]
        if args.max_bus_routes is not None:
            refs = list(dict.fromkeys(bus["ref"].dropna()))
            keep_refs = set(refs[:args.max_bus_routes])
            bus = bus[bus["ref"].isin(keep_refs)]
        out = gpd.GeoDataFrame(pd.concat([other, bus], ignore_index=True),
                               crs=out.crs)
    return out


def add_basemap(ax, styling, zoom=None):
    if zoom is None:
        zoom = styling["basemap"]["zoom"]
    try:
        import contextily as cx
        provider = styling["basemap"]["provider"]
        src = cx.providers
        for part in provider.split("."):
            src = getattr(src, part)
        cx.add_basemap(ax, crs=TARGET_CRS, source=src, zoom=zoom,
                       attribution=False, zorder=0,
                       headers={"User-Agent": UA})
        return True
    except Exception as e:
        print(f"  [basemap] unavailable ({e}); using plain background", flush=True)
        ax.set_facecolor("#f2f2f0")
        return False


def shapeburst_mask(ax, area_3857, extent, styling, res=1600):
    fade_m = styling["mask"]["fade_distance_m"]
    max_alpha = styling["mask"]["max_alpha"]
    color = styling["mask"]["color"]
    x0, x1, y0, y1 = extent
    nx = res
    ny = max(1, int(round(res * (y1 - y0) / (x1 - x0))))
    pxw = (x1 - x0) / nx
    pxh = (y1 - y0) / ny
    transform = Affine.translation(x0, y1) * Affine.scale(pxw, -pxh)

    from rasterio.features import rasterize
    geom = area_3857.geometry.union_all()
    mask = rasterize([(geom, 1)], out_shape=(ny, nx), transform=transform,
                     fill=0, dtype=np.uint8, all_touched=False).astype(bool)
    dist_px = distance_transform_edt(~mask)
    dist_m = dist_px * pxw
    alpha = np.where(mask, 0.0, np.clip(dist_m / fade_m, 0.0, 1.0)) * max_alpha

    cmap = matplotlib.colors.to_rgba(color)
    rgba = np.zeros((ny, nx, 4))
    rgba[..., :3] = cmap[:3]
    rgba[..., 3] = alpha
    ax.imshow(rgba, extent=(x0, x1, y0, y1), origin="upper",
              aspect="equal", interpolation="bilinear", zorder=0.5)


def plot_lines(ax, lines, styling):
    zord = styling["zorder"]
    widths = styling["line_width"]
    alphas = styling["alpha"]
    outline_w = styling["outline_width"]
    outline_c = styling["outline_color"]
    for style in sorted(zord, key=lambda k: zord[k]):
        grp = lines[lines["style"] == style]
        if grp.empty:
            continue
        w = widths.get(style, 1.0)
        z = zord[style]
        # dark casing
        grp.plot(ax=ax, color=outline_c, linewidth=w + outline_w,
                 zorder=z, alpha=1.0, capstyle="round", joinstyle="round")
        # colour body
        body_colors = grp["colour_final"].tolist()
        grp.plot(ax=ax, color=body_colors, linewidth=w,
                 zorder=z + 0.1, alpha=alphas.get(style, 1.0),
                 capstyle="round", joinstyle="round")


RAIL_STYLES = ("metro", "s_tog", "light_rail", "regional")


def classify_station(r):
    """Style key for a station row, or None. Stations are pre-classified."""
    s = r.get("style")
    return s if s in RAIL_STYLES else None


def plot_stations(ax, stations, styling, active_styles):
    """Draw station markers. Input is pre-collapsed: one row per (name, style)."""
    if stations is None or stations.empty:
        return
    cfg = styling.get("stations", {})
    if not cfg.get("show", True):
        return
    marker_styles = set(cfg.get("styles", ["metro", "s_tog", "regional"]))
    mk = cfg.get("marker", {})
    sizes = mk.get("size", {})
    shape = mk.get("shape", "circle")
    fill = mk.get("fill", "white")
    edge = mk.get("edge", "#2b2b2b")
    lw = mk.get("linewidth", 0.8)
    marker = "o" if shape == "circle" else "s"

    pts_by_style = {}
    for _, r in stations.iterrows():
        s = classify_station(r)
        if s is None or s not in marker_styles or s not in active_styles:
            continue
        pts_by_style.setdefault(s, []).append((r.geometry.x, r.geometry.y))

    zord = styling["zorder"]
    top_z = max(zord.values()) + 1  # all markers above all lines
    for s in sorted(pts_by_style, key=lambda k: zord.get(k, 0)):
        pts = pts_by_style[s]
        sz = sizes.get(s, 3.0)
        ax.scatter([p[0] for p in pts], [p[1] for p in pts],
                   s=sz ** 2, marker=marker, c=fill, edgecolors=edge,
                   linewidths=lw, zorder=top_z, alpha=1.0)


def rail_station_names(stations):
    """Set of names of rail stations (metro/s_tog/light_rail/regional)."""
    if stations is None or stations.empty:
        return set()
    names = set()
    for _, r in stations.iterrows():
        if classify_station(r) is not None:
            name = r.get("name")
            if isinstance(name, str) and name.strip():
                names.add(name)
    return names


def plot_bus_stops(ax, stops, styling, active_styles, rail_names=None):
    """Draw small markers for bus stops (pre-collapsed: one point per name).

    A bus stop whose name exactly matches a rail station is skipped — the
    station marker represents it.
    """
    if stops is None or stops.empty or not ({"bus", "ferry"} & active_styles):
        return
    cfg = styling.get("stations", {})
    if not cfg.get("show", True) or "bus" not in set(cfg.get("styles", [])):
        return
    bus = stops[stops["style"].isin(["bus", "ferry"])]
    if bus.empty:
        return

    rail_names = rail_names or set()
    mk = cfg.get("marker", {})
    sz = mk.get("size", {}).get("bus", 1.5)
    fill = mk.get("fill", "white")
    edge = mk.get("edge", "#2b2b2b")
    lw = mk.get("linewidth", 0.8)
    marker = "o" if mk.get("shape", "circle") == "circle" else "s"
    top_z = max(styling["zorder"].values()) + 1

    pts = [(r.geometry.x, r.geometry.y)
           for _, r in bus.iterrows()
           if not (isinstance(r.get("name"), str) and r.get("name") in rail_names)]
    ax.scatter([p[0] for p in pts], [p[1] for p in pts],
               s=sz ** 2, marker=marker, c=fill, edgecolors=edge,
               linewidths=lw, zorder=top_z, alpha=0.8)


def label_stations(ax, stations, styling, active_styles):
    if stations is None or stations.empty:
        return
    cfg = styling.get("labels", {})
    if not cfg.get("show", False):
        return
    label_styles = set(cfg.get("styles", ["metro", "s_tog"])) & active_styles
    fmin = cfg.get("min_fontsize", 5)
    fmax = cfg.get("max_fontsize", 9)

    # stations are unique per (name, style); dedupe by name for labelling
    seen = set()
    uniq = []
    for _, r in stations.iterrows():
        s = classify_station(r)
        if s not in label_styles:
            continue
        name = r.get("name")
        if not isinstance(name, str) or not name.strip() or name in seen:
            continue
        seen.add(name)
        uniq.append((r.geometry.x, r.geometry.y, name, s == "metro"))

    if not uniq:
        return
    texts = []
    for x, y, name, is_m in uniq:
        fs = fmax if is_m else fmin
        bbox = dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.7)
        t = ax.text(x, y, name, fontsize=fs, color="#111111",
                    zorder=10, ha="center", va="center", bbox=bbox)
        texts.append(t)
    try:
        adjust_text(texts, ax=ax, expand_points=(1.4, 1.6),
                    force_text=(0.4, 0.6), lim=300,
                    arrowprops=dict(arrowstyle="-", color="#888", lw=0.4))
    except Exception as e:
        print(f"  [labels] adjustText failed ({e}); labels may overlap", flush=True)


def main():
    args = load_args()
    styling = yaml.safe_load((CONFIG / "styling.yaml").read_text())

    area = gpd.read_file(PROCESSED / "area.gpkg").to_crs(TARGET_CRS)
    lines = gpd.read_file(PROCESSED / "master.gpkg", layer="lines").to_crs(TARGET_CRS)
    try:
        stations = gpd.read_file(PROCESSED / "master.gpkg", layer="stations").to_crs(TARGET_CRS)
    except Exception:
        stations = None
    try:
        stops = gpd.read_file(PROCESSED / "master.gpkg", layer="stops").to_crs(TARGET_CRS)
    except Exception:
        stops = None

    lines = filter_lines(lines, args)
    print(f"rendering {len(lines)} line features", flush=True)

    bounds = area.total_bounds  # (minx, miny, maxx, maxy)
    dx = bounds[2] - bounds[0]
    dy = bounds[3] - bounds[1]
    aspect = dx / dy
    pad = 0.12
    extent = (bounds[0] - pad * dx, bounds[2] + pad * dx,
              bounds[1] - pad * dy, bounds[3] + pad * dy)

    fig_w, fig_h = parse_size(args.size, aspect)
    fig, ax = plt.subplots(figsize=(fig_w, fig_h), dpi=args.dpi)
    ax.set_xlim(extent[0], extent[1])
    ax.set_ylim(extent[2], extent[3])
    ax.set_aspect("equal")
    ax.axis("off")

    # basemap
    if not args.no_basemap:
        add_basemap(ax, styling)

    # fade mask outside area
    shapeburst_mask(ax, area, extent, styling)

    # lines
    plot_lines(ax, lines, styling)

    # active styles (which mode layers are actually drawn)
    active_styles = set(lines["style"].unique())

    # station markers (symbols, no text)
    if not args.no_stops:
        plot_stations(ax, stations, styling, active_styles)
        rail_names = rail_station_names(stations)
        plot_bus_stops(ax, stops, styling, active_styles, rail_names=rail_names)

    # station labels (opt-in via --labels)
    if args.labels:
        styling["labels"]["show"] = True
    if styling.get("labels", {}).get("show", False):
        label_stations(ax, stations, styling, active_styles)

    # title + attribution
    ax.text(0.5, 0.985, "Kurragömma Köpenhamn",
            transform=ax.transAxes, ha="center", va="top",
            fontsize=fig_w * 1.1, fontweight="bold", color="#222",
            bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.7))
    ax.text(0.01, 0.01,
            "Transit: Rejseplanen GTFS · Area: © OpenStreetMap contributors (ODbL) · Base: Esri, HERE",
            transform=ax.transAxes, ha="left", va="bottom",
            fontsize=max(5, fig_w * 0.5), color="#666")

    out_dir = OUTPUT
    out_dir.mkdir(parents=True, exist_ok=True)
    if args.out:
        out_path = Path(args.out)
        if not out_path.is_absolute():
            out_path = out_dir / out_path
    else:
        out_path = out_dir / f"copenhagen.{args.format}"
    out_path.parent.mkdir(parents=True, exist_ok=True)

    plt.savefig(out_path, dpi=args.dpi, bbox_inches="tight",
                pad_inches=0.2, facecolor="white")
    plt.close(fig)
    print(f"wrote {out_path} ({args.format}, {fig_w}x{fig_h}in @ {args.dpi}dpi)")


if __name__ == "__main__":
    main()