"""TECH-027 phase 2: CP ride geometry from OpenStreetMap rails.

CP's GTFS has no shapes.txt. For every pair of consecutive CP calls (stop_times), this finds
the rail path between the tracks at the two stations in OSM and writes it, simplified to 3 m,
into a small SQLite file the backend ships as a resource (cp-rails.sqlite).

Graph: OSM ways railway=rail|narrow_gauge (light_rail, subway and tram are other networks).
Search is edge-based so a path cannot take a sharp turn at a switch or a diamond crossing
(MAX_TURN_DEG); sidings, yards and industrial track cost more so a path does not shortcut
through a yard. A pair is stored once, as shape id "<a>><b>" with a < b; the backend reverses it for b → a.

    osmium tags-filter portugal.osm.pbf w/railway=rail,narrow_gauge -o rails.osm.pbf
    osmium add-locations-to-ways rails.osm.pbf -f opl -o rails.opl
    python3 -I tools/cp-rails/build.py --opl rails.opl --gtfs CP.zip --osm-date YYYY-MM-DD \
        --out backend/src/main/resources/routing/cp-rails.sqlite --report /tmp/cp-rails-report.json \
        [--geojson /tmp/cp-rails.geojson]

Python 3 only, no dependencies. Data: © OpenStreetMap contributors (ODbL 1.0).
"""
import argparse, csv, heapq, io, json, math, os, sqlite3, sys, zipfile
from collections import defaultdict, Counter

FORMAT = "v1"
SIMPLIFY_METERS = 3.0
MAX_TURN_DEG = 50.0          # deviation between consecutive track segments a train can take
SNAP_RADIUS_M = (250.0, 600.0)  # tracks searched around a station, wider if none
SNAP_COST = 2.0              # metres of path per metre a track lies away from the station
FACTOR = {"siding": 2.0, "spur": 2.0, "yard": 3.0, "crossover": 1.2}
USAGE_FACTOR = {"industrial": 3.0, "military": 3.0, "test": 3.0}
M_LAT = 110_540.0
M_LON = 111_320.0
COS0 = math.cos(math.radians(39.5))


def unescape(s):
    # OPL escapes anything outside a safe set as %hex%
    out, i = [], 0
    while i < len(s):
        if s[i] == "%":
            j = s.index("%", i + 1)
            out.append(chr(int(s[i + 1:j], 16))); i = j + 1
        else:
            out.append(s[i]); i += 1
    return "".join(out)


def read_rails(path):
    """Ways from add-locations-to-ways OPL: [(factor, [(osm_node_id, lat, lon), ...]), ...]."""
    ways = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            if not line.startswith("w"):
                continue
            tags, nodes = {}, []
            for field in line.rstrip("\n").split(" "):
                if field.startswith("T") and len(field) > 1:
                    for kv in field[1:].split(","):
                        if "=" in kv:
                            k, v = kv.split("=", 1); tags[unescape(k)] = unescape(v)
                elif field.startswith("N"):
                    for ref in field[1:].split(","):
                        if "x" not in ref:
                            continue  # node without a location (outside the extract)
                        nid, rest = ref[1:].split("x", 1)
                        lon, lat = rest.split("y", 1)
                        nodes.append((int(nid), float(lat), float(lon)))
            if tags.get("railway") not in ("rail", "narrow_gauge") or len(nodes) < 2:
                continue
            if tags.get("disused") == "yes" or tags.get("abandoned") == "yes":
                continue
            factor = max(FACTOR.get(tags.get("service", ""), 1.0), USAGE_FACTOR.get(tags.get("usage", ""), 1.0))
            ways.append((factor, nodes))
    return ways


def meters(lat1, lon1, lat2, lon2):
    k = math.cos(math.radians((lat1 + lat2) / 2))
    return math.hypot((lon2 - lon1) * M_LON * k, (lat2 - lat1) * M_LAT)


class Graph:
    def __init__(self, ways):
        idx = {}
        self.lat, self.lon = [], []
        self.tail, self.head, self.cost, self.length = [], [], [], []
        self.out = []
        for factor, nodes in ways:
            prev = None
            for nid, lat, lon in nodes:
                i = idx.get(nid)
                if i is None:
                    i = idx[nid] = len(self.lat)
                    self.lat.append(lat); self.lon.append(lon); self.out.append([])
                if prev is not None and prev != i:
                    d = meters(self.lat[prev], self.lon[prev], lat, lon)
                    for a, b in ((prev, i), (i, prev)):
                        e = len(self.tail)
                        self.tail.append(a); self.head.append(b); self.length.append(d); self.cost.append(d * factor)
                        self.out[a].append(e)
                prev = i
        self.grid = defaultdict(list)
        for i in range(len(self.lat)):
            self.grid[self.cell(self.lat[i], self.lon[i])].append(i)

    @staticmethod
    def cell(lat, lon):
        return (int(lat * M_LAT // 500), int(lon * M_LON * COS0 // 500))

    def near(self, lat, lon, radius):
        cy, cx = self.cell(lat, lon)
        r = int(radius // 500) + 1
        found = []
        for dy in range(-r, r + 1):
            for dx in range(-r, r + 1):
                for i in self.grid.get((cy + dy, cx + dx), ()):
                    d = meters(lat, lon, self.lat[i], self.lon[i])
                    if d <= radius:
                        found.append((i, d))
        return found

    def bearing(self, e):
        a, b = self.tail[e], self.head[e]
        return math.atan2((self.lat[b] - self.lat[a]) * M_LAT, (self.lon[b] - self.lon[a]) * M_LON * COS0)

    def turn_ok(self, e1, e2):
        if self.head[e2] == self.tail[e1]:
            return False  # no reversing onto the segment just travelled
        d = abs(self.bearing(e2) - self.bearing(e1))
        d = min(d, 2 * math.pi - d)
        return math.degrees(d) <= MAX_TURN_DEG

    def path(self, a, b):
        """Cheapest rail path between the tracks around stations a and b: (node list, snap_a, snap_b) or None."""
        for radius in SNAP_RADIUS_M:
            src, dst = self.near(a[0], a[1], radius), self.near(b[0], b[1], radius)
            if src and dst:
                break
        else:
            return None
        goal = {i: d for i, d in dst}
        heap, g, parent = [], {}, {}
        def h(n):
            return max(0.0, meters(self.lat[n], self.lon[n], b[0], b[1]) - radius)
        for s, d in src:
            for e in self.out[s]:
                c = SNAP_COST * d + self.cost[e]
                if c < g.get(e, math.inf):
                    g[e] = c; parent[e] = None
                    heapq.heappush(heap, (c + h(self.head[e]), c, e))
        best, best_e = math.inf, None
        while heap:
            f, c, e = heapq.heappop(heap)
            if f >= best:
                break
            if c > g.get(e, math.inf):
                continue
            n = self.head[e]
            if n in goal and c + SNAP_COST * goal[n] < best:
                best, best_e = c + SNAP_COST * goal[n], e
            for e2 in self.out[n]:
                if not self.turn_ok(e, e2):
                    continue
                c2 = c + self.cost[e2]
                if c2 < g.get(e2, math.inf):
                    g[e2] = c2; parent[e2] = e
                    heapq.heappush(heap, (c2 + h(self.head[e2]), c2, e2))
        if best_e is None:
            return None
        edges = []
        e = best_e
        while e is not None:
            edges.append(e); e = parent[e]
        edges.reverse()
        nodes = [self.tail[edges[0]]] + [self.head[e] for e in edges]
        snap_a = meters(a[0], a[1], self.lat[nodes[0]], self.lon[nodes[0]])
        snap_b = meters(b[0], b[1], self.lat[nodes[-1]], self.lon[nodes[-1]])
        return nodes, snap_a, snap_b


def simplify(xs, ys, tol=SIMPLIFY_METERS):
    """Douglas–Peucker on metre coordinates; same rule as ShapeGeometry.simplify."""
    n = len(xs)
    if n <= 2:
        return list(range(n))
    keep = [False] * n; keep[0] = keep[-1] = True
    stack = [(0, n - 1)]
    while stack:
        i, j = stack.pop()
        ax, ay, bx, by = xs[i], ys[i], xs[j], ys[j]
        dx, dy = bx - ax, by - ay
        L2 = dx * dx + dy * dy
        worst, wk = -1.0, -1
        for k in range(i + 1, j):
            px, py = xs[k] - ax, ys[k] - ay
            if L2 == 0:
                d = math.hypot(px, py)
            else:
                t = max(0.0, min(1.0, (px * dx + py * dy) / L2))
                d = math.hypot(px - t * dx, py - t * dy)
            if d > worst:
                worst, wk = d, k
        if worst > tol:
            keep[wk] = True
            stack.append((i, wk)); stack.append((wk, j))
    return [k for k in range(n) if keep[k]]


def zigzag(v):
    return 2 * v if v >= 0 else -2 * v - 1  # = Kotlin (v shl 1) xor (v shr 31) for 32-bit v


def encode(lat, lon):
    out = bytearray(); plat = plon = 0
    for a, o in zip(lat, lon):
        for d in (a - plat, o - plon):
            v = zigzag(d)
            while v & ~0x7F:
                out.append((v & 0x7F) | 0x80); v >>= 7
            out.append(v)
        plat, plon = a, o
    return bytes(out)


def read_gtfs(path):
    with zipfile.ZipFile(path) as z:
        def rows(name):
            return csv.DictReader(io.TextIOWrapper(z.open(name), encoding="utf-8-sig"), skipinitialspace=True)
        stops = {r["stop_id"]: (float(r["stop_lat"]), float(r["stop_lon"]), r["stop_name"]) for r in rows("stops.txt")}
        calls = defaultdict(list)
        for r in rows("stop_times.txt"):
            calls[r["trip_id"]].append((int(r["stop_sequence"]), r["stop_id"]))
    pairs = Counter()
    for seq in calls.values():
        seq.sort()
        for (_, s1), (_, s2) in zip(seq, seq[1:]):
            if s1 != s2:
                pairs[tuple(sorted((s1, s2)))] += 1
    return stops, pairs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--opl", required=True)
    ap.add_argument("--gtfs", required=True)
    ap.add_argument("--osm-date", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--report", required=True)
    ap.add_argument("--geojson")
    args = ap.parse_args()

    stops, pairs = read_gtfs(args.gtfs)
    graph = Graph(read_rails(args.opl))
    print(f"rails: {len(graph.lat)} nodes, {len(graph.tail) // 2} segments; CP: {len(stops)} stops, {len(pairs)} pairs", file=sys.stderr)

    rows, report, feats = [], [], []
    for k, ((s1, s2), trips) in enumerate(sorted(pairs.items())):
        a, b = stops[s1], stops[s2]
        straight = meters(a[0], a[1], b[0], b[1])
        found = graph.path(a, b)
        entry = {"from": s1, "to": s2, "fromName": a[2], "toName": b[2], "trips": trips, "straightM": round(straight)}
        if found is None:
            entry["miss"] = "NO_TRACK"
            report.append(entry); continue
        nodes, snap_a, snap_b = found
        lat = [graph.lat[n] for n in nodes]; lon = [graph.lon[n] for n in nodes]
        length = sum(meters(lat[i], lon[i], lat[i + 1], lon[i + 1]) for i in range(len(lat) - 1))
        xs = [o * M_LON * COS0 for o in lon]; ys = [a_ * M_LAT for a_ in lat]
        keep = simplify(xs, ys)
        mlat = [round(lat[i] * 1e6) for i in keep]; mlon = [round(lon[i] * 1e6) for i in keep]
        entry.update(pathM=round(length), ratio=round(length / straight, 3) if straight else None,
                     snapFromM=round(snap_a), snapToM=round(snap_b), points=len(keep))
        report.append(entry)
        rows.append((s1, s2, len(keep), encode(mlat, mlon)))
        if args.geojson:
            feats.append({"type": "Feature", "properties": {k_: v for k_, v in entry.items()},
                          "geometry": {"type": "LineString", "coordinates": [[o / 1e6, a_ / 1e6] for a_, o in zip(mlat, mlon)]}})
        if k % 100 == 0:
            print(f"  {k}/{len(pairs)}", file=sys.stderr)

    tmp = args.out + ".tmp"
    if os.path.exists(tmp):
        os.remove(tmp)
    db = sqlite3.connect(tmp)
    db.execute("PRAGMA journal_mode=OFF")
    # Same table as the feed shape files (TransitShapeStore), so one reader serves both.
    db.execute("CREATE TABLE shapes (id TEXT PRIMARY KEY, n INTEGER NOT NULL, pts BLOB NOT NULL) WITHOUT ROWID")
    db.execute("CREATE TABLE meta (key TEXT PRIMARY KEY, value TEXT NOT NULL) WITHOUT ROWID")
    db.executemany("INSERT INTO shapes VALUES (?, ?, ?)", [(f"{a}>{b}", n, pts) for a, b, n, pts in rows])
    db.executemany("INSERT INTO meta VALUES (?, ?)", [
        ("format", FORMAT), ("osm_date", args.osm_date), ("gtfs", os.path.basename(args.gtfs)),
        ("license", "Rail geometry © OpenStreetMap contributors, ODbL 1.0"),
        ("simplify_m", str(SIMPLIFY_METERS)), ("max_turn_deg", str(MAX_TURN_DEG))])
    db.commit(); db.execute("VACUUM"); db.close()
    os.replace(tmp, args.out)

    json.dump(report, open(args.report, "w"), ensure_ascii=False, indent=1)
    if args.geojson:
        json.dump({"type": "FeatureCollection", "features": feats}, open(args.geojson, "w"))
    found = [r for r in report if "miss" not in r]
    ratios = sorted(r["ratio"] for r in found if r["ratio"])
    print(f"pairs {len(report)}: found {len(found)}, no track {len(report) - len(found)}; "
          f"points {sum(r['points'] for r in found)}; file {os.path.getsize(args.out) / 1024:.0f} KB", file=sys.stderr)
    if ratios:
        print("path/straight p50 %.2f p90 %.2f p99 %.2f max %.2f" % (ratios[len(ratios) // 2], ratios[int(len(ratios) * .9)],
              ratios[min(len(ratios) - 1, int(len(ratios) * .99))], ratios[-1]), file=sys.stderr)


if __name__ == "__main__":
    main()