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

Reads data/processed/master.gpkg (layers: lines, stations) + area.gpkg +
config/styling.yaml, applies CLI filters, and composes a printable map:
  - Carto Positron (no labels) 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

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 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]), crs=out.crs)
    return out


def pd_concat(frames):
    import pandas as pd
    return pd.concat(frames, ignore_index=True)


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")


def classify_station(r):
    """Return the style key for a station row, or None."""
    if r.get("subway") == "yes" or r.get("station") == "subway":
        return "metro"
    if r.get("light_rail") == "yes" or r.get("station") == "light_rail":
        return "s_tog"
    if r.get("railway") == "station" or r.get("train") == "yes":
        return "regional"
    return None


def cluster_stops(coords, threshold_m):
    """Cluster (x, y) coordinates within threshold_m; return list of centroids."""
    from scipy.cluster.hierarchy import fcluster, linkage
    from scipy.spatial.distance import pdist

    coords = np.asarray(coords)
    if len(coords) <= 1:
        return coords.tolist() if len(coords) else []

    dists = pdist(coords)
    links = linkage(dists, method="single")
    labels = fcluster(links, t=threshold_m, criterion="distance")

    centroids = []
    for label in set(labels):
        members = coords[labels == label]
        centroids.append(members.mean(axis=0).tolist())
    return centroids


def plot_stations(ax, stations, styling, active_styles):
    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"]))
    cluster_cfg = cfg.get("stop_cluster_m", {})
    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"

    # group coordinates by style, then by name for clustering
    pts_by_style = {}
    names_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
        coord = (r.geometry.x, r.geometry.y)
        name = r.get("name") if isinstance(r.get("name"), str) else None
        pts_by_style.setdefault(s, []).append(coord)
        names_by_style.setdefault(s, {}).setdefault(name, []).append(coord)

    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)):
        threshold = cluster_cfg.get(s, 50) if isinstance(cluster_cfg, dict) else cluster_cfg
        # cluster same-name stations to one marker; keep unnamed as-is
        centroids = []
        for name, coords in names_by_style[s].items():
            if name and len(coords) > 1:
                centroids.extend(cluster_stops(coords, threshold))
            else:
                centroids.extend(coords)
        sz = sizes.get(s, 3.0)
        ax.scatter([p[0] for p in centroids], [p[1] for p in centroids],
                   s=sz ** 2, marker=marker, c=fill, edgecolors=edge,
                   linewidths=lw, zorder=top_z, alpha=1.0)


def rail_station_names(stations):
    """Return the set of names for rail stations (metro/s_tog/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 from the stops layer.

    Bus stops whose name matches a rail station are skipped — the rail
    station marker represents that stop.
    """
    if stops is None or stops.empty or "bus" not in 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"] == "bus"]
    if bus.empty:
        return

    if rail_names is None:
        rail_names = set()

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

    # cluster same-name bus stops within threshold (opposite sides of road)
    # skip names that match a rail station (merged into station marker)
    centroids = []
    skipped = 0
    for name, grp in bus.groupby("name"):
        if name and name in rail_names:
            skipped += len(grp)
            continue
        coords = [(r.geometry.x, r.geometry.y) for _, r in grp.iterrows()]
        if name and len(coords) > 1:
            centroids.extend(cluster_stops(coords, cluster_m))
        else:
            centroids.extend(coords)

    ax.scatter([p[0] for p in centroids], [p[1] for p in centroids],
               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)

    def is_metro(r):
        return classify_station(r) == "metro"

    def is_stog(r):
        return classify_station(r) == "s_tog"

    pts = []
    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():
            continue
        pts.append((r.geometry.x, r.geometry.y, name, s == "metro"))

    if not pts:
        return
    # dedupe by name (keep first location)
    seen = {}
    uniq = []
    for x, y, name, is_m in pts:
        if name in seen:
            continue
        seen[name] = True
        uniq.append((x, y, name, is_m))

    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,
            "Data: © 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()