#!/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
  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

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("--no-labels", action="store_true", help="skip station labels")
    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)
        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 label_stations(ax, stations, styling, area_3857):
    if stations is None or stations.empty:
        return
    label_styles = set(styling["labels"]["styles"])
    fmin = styling["labels"]["min_fontsize"]
    fmax = styling["labels"]["max_fontsize"]

    def is_metro(r):
        return r.get("subway") == "yes" or r.get("station") == "subway"

    def is_stog(r):
        return r.get("light_rail") == "yes" or r.get("station") == "light_rail"

    pts = []
    for _, r in stations.iterrows():
        is_m = is_metro(r)
        is_s = is_stog(r)
        if "metro" in label_styles and is_m:
            pass
        elif "s_tog" in label_styles and is_s:
            pass
        else:
            continue
        name = r.get("name")
        if not isinstance(name, str) or not name.strip():
            continue
        pts.append((r.geometry.x, r.geometry.y, name, is_m))

    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

    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)

    # labels
    if not args.no_labels and styling["labels"]["show"]:
        label_stations(ax, stations, styling, area)

    # title + attribution
    ax.text(0.5, 0.985, "Copenhagen — Public Transit",
            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: © CARTO",
            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()