#!/usr/bin/env python3
"""Bereken autoroutes uit Nederlandse PC6-postcodes met lokale OSRM."""

import argparse
from pathlib import Path

import geopandas as gpd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.ticker import PercentFormatter
import numpy as np
import pandas as pd
import requests
from shapely.geometry import LineString

PC6 = r"[1-9][0-9]{3}[A-Z]{2}"

EXAMPLES = [
    ("1791VG", "Texel", "#007d68", "--"),
    ("8881CE", "Terschelling", "#a32862", "-"),
    ("9161BZ", "Ameland", "#6c46a4", "-."),
    ("6291KJ", "Vaals", "#a24d00", ":"),
]

def normalize(series):
    return series.astype("string").str.upper().str.replace(r"\s+", "", regex=True)


def lookup_points(args, codes):
    if args.cbs:
        # Codes have already been validated. Read only the relevant PC6 polygons.
        quoted = ", ".join(f"'{code}'" for code in codes)
        areas = gpd.read_file(
            args.cbs, layer=args.cbs_laag, engine="pyogrio",
            columns=["postcode6"], where=f"postcode6 IN ({quoted})",
        ).to_crs(28992)
        points = areas.dissolve(by="postcode6").geometry.representative_point().to_crs(4326)
        return pd.DataFrame({"lon": points.x, "lat": points.y})
    points = pd.read_csv(args.punten, dtype={"postcode": "string"})
    points["postcode"] = normalize(points.postcode)
    if points.postcode.duplicated().any():
        raise ValueError("Gebruik één punt per postcode in de punten-CSV.")
    return points.set_index("postcode")[["lon", "lat"]]


def route(session, origin, destination):
    coordinates = ";".join(f"{lon},{lat}" for lon, lat in (origin, destination))
    response = session.get(
        f"http://127.0.0.1:5000/route/v1/driving/{coordinates}",
        params={"overview": "simplified", "geometries": "geojson"},
        timeout=45,
    )
    payload = response.json()
    if payload.get("code") in {"NoRoute", "NoSegment"}:
        return None
    response.raise_for_status()
    return payload["routes"][0]


def plot_maps(data, lines, destination, boundaries_path, output):
    if boundaries_path and boundaries_path.is_file():
        boundaries = gpd.read_file(boundaries_path).to_crs(28992)
    else:
        boundaries = None
    points = gpd.GeoSeries(gpd.points_from_xy(data.lon, data.lat), crs=4326).to_crs(28992)
    target = gpd.GeoSeries(gpd.points_from_xy([destination[0]], [destination[1]]), crs=4326).to_crs(28992).iloc[0]
    geometry = gpd.GeoSeries(lines, crs=4326).to_crs(28992)
    extent = boundaries.total_bounds if boundaries is not None else geometry.total_bounds
    palette = ["#c4d722", "#27a384", "#347995", "#492372"]
    travel_classes = np.digitize(data.retour_min, [60, 120, 180], right=True)

    for mode, title, subtitle in [
        ("herkomst", "Herkomst van de voorbeeldgroep", "Stipgrootte volgt het aantal per postcode"),
        ("retourtijd", "Retourtijd per postcode", "Heen- en terugrit afzonderlijk berekend"),
        ("routes", "Autoroutes naar de bestemming", "Heenroutes · lijndikte volgt het aantal per postcode"),
    ]:
        fig, ax = plt.subplots(figsize=(7.2, 7.6))
        fig.subplots_adjust(left=.08, right=.96, top=.84, bottom=.13)
        fig.text(.08, .95, title, fontsize=19, weight="bold")
        fig.text(.08, .90, subtitle, fontsize=10, color="#526675")
        if boundaries is not None:
            boundaries.boundary.plot(ax=ax, color="#bdc9cf", linewidth=.45)
        if mode == "routes":
            geometry.plot(ax=ax, color="#004488", alpha=.23,
                          linewidth=.35 + .065*data.aantal.to_numpy())
        elif mode == "retourtijd":
            dots = ax.scatter(points.x, points.y, c=travel_classes, s=35,
                              cmap=ListedColormap(palette),
                              norm=BoundaryNorm(np.arange(-.5, 4.5), 4),
                              edgecolor="white", linewidth=.4, zorder=4)
            bar = fig.colorbar(dots, ax=ax, orientation="horizontal", shrink=.8, pad=.035,
                               fraction=.035, ticks=[0, 1, 2, 3])
            bar.ax.set_xticklabels(["≤1 uur", ">1–2 uur", ">2–3 uur", ">3 uur"])
        else:
            ax.scatter(points.x, points.y, s=10+3*data.aantal.to_numpy(),
                       color="#004488", edgecolor="white", linewidth=.5, zorder=4)
        ax.scatter(target.x, target.y, marker="*", s=130, color="#bb5566",
                   edgecolor="white", linewidth=.8, zorder=5)
        ax.set(xlim=(extent[0]-7000, extent[2]+7000), ylim=(extent[1]-7000, extent[3]+7000),
               aspect="equal")
        ax.set_axis_off()
        credit = "Postcodegeometrie CBS / Esri Nederland · Wegen © OpenStreetMap contributors"
        if boundaries is not None:
            credit += "\nGemeentegrenzen CBS / Kadaster"
        fig.text(.08, .03, credit, fontsize=8, color="#526675")
        fig.savefig(output / f"{mode}.png", dpi=170, facecolor="white")
        plt.close(fig)


    plot_example_routes(data, geometry, target, boundaries, output)


def plot_example_routes(data, geometry, target, boundaries, output):
    """Teken dezelfde vier voorbeeldroutes als in de histogrammen."""
    by_postcode = dict(zip(data.postcode, geometry))
    examples = [item for item in EXAMPLES if item[0] in by_postcode]
    if not examples:
        return
    fig, ax = plt.subplots(figsize=(7.8, 8.2))
    fig.subplots_adjust(left=.07, right=.96, top=.79, bottom=.10)
    if boundaries is not None:
        boundaries.boundary.plot(ax=ax, color="#c6d0d5", linewidth=.45)
    for postcode, name, color, style in examples:
        x, y = by_postcode[postcode].xy
        ax.plot(x, y, color=color, linestyle=style, linewidth=2, label=name)
        ax.scatter(x[0], y[0], color=color, s=33, edgecolor="white", linewidth=.6, zorder=5)
    ax.scatter(target.x, target.y, marker="*", s=140, color="#243746",
               edgecolor="white", linewidth=.8, zorder=6)
    ax.annotate("Bestemming", (target.x, target.y), xytext=(8, -9),
                textcoords="offset points", fontsize=9)
    ax.set_aspect("equal")
    ax.set_axis_off()
    handles, labels = ax.get_legend_handles_labels()
    fig.legend(handles, labels, loc="upper left", bbox_to_anchor=(.074,.865),
               ncol=2, frameon=False, fontsize=10)
    fig.text(.08, .95, "Vier voorbeeldroutes op de kaart", fontsize=19, weight="bold")
    fig.text(.08, .905, "Heenroutes · dezelfde kleuren en lijnstijlen als in de verdelingen",
             fontsize=10, color="#526675")
    credit = "Postcodegeometrie CBS / Esri Nederland · Wegen © OpenStreetMap contributors"
    if boundaries is not None:
        credit += "\nGemeentegrenzen CBS / Kadaster"
    fig.text(.08, .035, credit, fontsize=8, color="#526675")
    fig.savefig(output / "voorbeeldroutes.png", dpi=170, facecolor="white")
    plt.close(fig)


def plot_distribution(data, output):
    by_postcode = data.set_index("postcode")
    for column, step, unit, title, xlabel, filename in [
        ("heen_km", 20, "km", "Verdeling van de enkele reisafstand",
         "Afstand heen (km)", "afstand-verdeling.png"),
        ("retour_min", 30, "min", "Verdeling van de retourrijtijd",
         "Tijd heen en terug (minuten)", "verdeling.png"),
    ]:
        values = data[column].to_numpy()
        upper = max(step, int(np.ceil(values.max()/step))*step)
        edges = np.arange(0, upper+step, step)
        # Exacte bovengrenzen horen bij de lagere klasse.
        bins = np.digitize(values, edges[1:-1], right=True)
        counts = np.bincount(bins, weights=data.aantal, minlength=len(edges)-1)
        assert counts.sum() == data.aantal.sum()
        shares = counts/counts.sum()
        fig, ax = plt.subplots(figsize=(9, 5.8))
        fig.subplots_adjust(left=.12, right=.96, top=.67, bottom=.17)
        fig.text(.12, .95, title, fontsize=19, weight="bold")
        fig.text(.12, .90, f"Klassen van {step} {unit} · gewogen naar aantal", fontsize=10,
                 color="#526675")
        ax.bar(edges[:-1]+step/2, shares, width=step*.86, color="#b5cbdc")
        for postcode, name, color, style in EXAMPLES:
            if postcode in by_postcode.index:
                value = float(by_postcode.loc[postcode, column])
                ax.axvline(value, color=color, linestyle=style, linewidth=2,
                           label=f"{name} · {value:.0f} {unit}", zorder=4)
        handles, labels = ax.get_legend_handles_labels()
        if handles:
            fig.legend(handles, labels, loc="upper left", bbox_to_anchor=(.115,.865),
                       ncol=2, frameon=False, fontsize=10)
        ax.set_xlim(0, upper)
        ax.set_xticks(np.arange(0, upper+1, step*2))
        ax.yaxis.set_major_formatter(PercentFormatter(1))
        ax.set_ylim(0, shares.max()*1.13)
        ax.set_xlabel(xlabel)
        ax.set_ylabel("Aandeel van de voorbeeldgroep")
        ax.spines[["top", "right"]].set_visible(False)
        ax.set_axisbelow(True)
        ax.grid(axis="y", color="#e0e6ea")
        fig.text(.12, .035, "OSRM · Wegen © OpenStreetMap contributors", fontsize=8,
                 color="#526675")
        fig.savefig(output / filename, dpi=170, facecolor="white")
        plt.close(fig)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=Path("postcodes.csv"))
    parser.add_argument("--punten", type=Path, default=Path("postcodepunten.csv"))
    parser.add_argument("--cbs", type=Path, help="PC6 GeoPackage in plaats van de punten-CSV")
    parser.add_argument("--cbs-laag", default="cbs_pc6_2025")
    parser.add_argument("--bestemming", default="1012JS")
    parser.add_argument("--grenzen", type=Path, default=Path("grenzen.geojson"))
    parser.add_argument("--output", type=Path, default=Path("uitvoer"))
    args = parser.parse_args()

    data = pd.read_csv(args.input, dtype={"postcode": "string"})
    data["postcode"] = normalize(data.postcode)
    if data.empty or not data.postcode.str.fullmatch(PC6).fillna(False).all():
        parser.error("Gebruik een lijst met geldige PC6-codes.")
    if "aantal" not in data:
        data["aantal"] = 1
    data["aantal"] = pd.to_numeric(data.aantal, errors="coerce")
    if not (data.aantal.gt(0) & data.aantal.mod(1).eq(0)).all():
        parser.error("Aantallen moeten positieve gehele getallen zijn.")
    data = data.groupby("postcode", as_index=False).aantal.sum()
    target = normalize(pd.Series([args.bestemming])).iloc[0]
    if not pd.Series([target]).str.fullmatch(PC6).iloc[0]:
        parser.error("De bestemming moet een PC6-code zijn.")
    points = lookup_points(args, sorted(set(data.postcode) | {target}))
    if target not in points.index or not np.isfinite(points.loc[target]).all():
        parser.error("Geen geldig postcodepunt voor de bestemming.")
    destination = tuple(points.loc[target])

    rows, lines = [], []
    with requests.Session() as session:
        session.trust_env = False
        for item in data.itertuples(index=False):
            result = dict(postcode=item.postcode, aantal=int(item.aantal), status="geen postcodepunt",
                          lon=None, lat=None, heen_km=None, heen_min=None, terug_min=None, retour_min=None)
            if item.postcode in points.index and np.isfinite(points.loc[item.postcode]).all():
                origin = tuple(points.loc[item.postcode])
                result.update(lon=origin[0], lat=origin[1], status="geen route")
                try:
                    outward = route(session, origin, destination)
                    inward = route(session, destination, origin) if outward else None
                except requests.RequestException as exc:
                    raise SystemExit(f"OSRM-aanvraag mislukt. Controleer de lokale server. {exc}") from exc
                if outward and inward:
                    result.update(status="ok", heen_km=outward["distance"]/1000,
                                  heen_min=outward["duration"]/60, terug_min=inward["duration"]/60,
                                  retour_min=(outward["duration"]+inward["duration"])/60)
                    lines.append(LineString(outward["geometry"]["coordinates"]))
            rows.append(result)

    args.output.mkdir(parents=True, exist_ok=True)
    result = pd.DataFrame(rows)
    result.to_csv(args.output / "reisafstanden.csv", index=False, encoding="utf-8-sig")
    valid = result[result.status == "ok"]
    print(f"Dekking: {len(valid)}/{len(result)} postcodes, {valid.aantal.sum()}/{data.aantal.sum()} personen")
    if valid.empty:
        raise SystemExit("Geen volledige routes. Bekijk de statuskolom in reisafstanden.csv.")
    for column in ["heen_km", "retour_min"]:
        print(f"{column}: gewogen gemiddelde {np.average(valid[column], weights=valid.aantal):.1f}")
    plot_maps(valid, lines, destination, args.grenzen, args.output)
    plot_distribution(valid, args.output)


if __name__ == "__main__":
    main()