#!/usr/bin/env python3
"""Build world maxspeed grid from planet PBF in 2 passes.

Pass 1: Stream PBF → collect node IDs from highway+maxspeed ways → bitset
Pass 2: Stream PBF → matching nodes → mmap sort → process ways → grid

Usage: python3 build_world.py planet-latest.osm.pbf grid-planet/
"""
import sys, os, json, gzip, struct, math, mmap, time
import numpy as np
import osmium

CELL = 0.5
SUB = 32

ZONE = {
    "none": None, "signals": None, "variable": None, "walk": 5, "living_street": 20,
    "urban": 50, "rural": 80, "trunk": 110, "motorway": 130,
}
CC = {
    "FR": (50,80), "BE": (50,90), "DE": (50,100), "NL": (50,100),
    "GB": (48,96), "ES": (50,90), "IT": (50,90), "CH": (50,80),
    "AT": (50,100), "LU": (50,90), "PT": (50,90), "PL": (50,90),
    "CZ": (50,90), "IE": (50,80), "DK": (50,80), "SE": (50,70),
    "NO": (50,80), "FI": (50,100), "RO": (50,90), "HU": (50,90),
    "HR": (50,90), "SI": (50,90), "SK": (50,90), "LT": (50,90),
    "LV": (50,90), "EE": (50,90), "BG": (50,90), "GR": (50,90),
    "UA": (50,90), "RU": (60,90), "TR": (50,90), "RS": (50,80),
    "MK": (50,80), "AL": (40,90), "BA": (50,80), "ME": (50,80),
    "IS": (50,90), "US": (40,105), "CA": (40,100), "MX": (40,90),
    "BR": (40,80), "AR": (40,110), "CL": (40,120), "CO": (40,90),
    "PE": (40,90), "JP": (40,60), "CN": (40,100), "IN": (40,80),
    "AU": (40,100), "NZ": (40,100), "ZA": (40,100), "EG": (50,90),
    "NG": (50,80), "KE": (50,80), "MA": (50,90), "DZ": (50,90),
    "TN": (50,90), "LY": (50,90),
}

def parse_speed(v):
    if not v:
        return None
    v = v.strip().lower()
    if v in ZONE:
        return ZONE[v]
    m = __import__("re").match(r"^([0-9]+)(?:\s*(mph|km/h|kmh))?\b", v)
    if m:
        n = int(m.group(1))
        if m.group(2) == "mph":
            return round(n * 1.609344)
        return n
    if ":" in v:
        cc, kind = v.split(":", 1)
        cc = cc.upper()
        if cc in CC and kind in ("urban", "rural"):
            return CC[cc][0 if kind == "urban" else 1]
        if cc in CC and kind == "motorway":
            return None
    return None

MAX_NODE_ID = 25_000_000_000  # 25 billion — covers all OSM node IDs with margin

class Pass1Collector(osmium.SimpleHandler):
    """Collect node IDs referenced by highway+maxspeed ways."""
    def __init__(self):
        super().__init__()
        bs_size = MAX_NODE_ID // 8 + 1
        sys.stderr.write(f"  Allocating bitset: {bs_size/1073741824:.1f} GB\n")
        self.bitset = bytearray(bs_size)
        self.way_count = 0

    def way(self, w):
        if "highway" not in w.tags:
            return
        speed = parse_speed(w.tags.get("maxspeed", ""))
        if speed is None:
            return
        self.way_count += 1
        for n in w.nodes:
            nid = n.ref
            if nid >= MAX_NODE_ID:
                continue
            byte_idx = nid >> 3
            bit_idx = nid & 7
            self.bitset[byte_idx] |= (1 << bit_idx)
        if self.way_count % 1_000_000 == 0:
            sys.stderr.write(f"  pass1: {self.way_count//1_000_000}M ways\n")

    def test(self, nid):
        if nid >= MAX_NODE_ID:
            return False
        return (self.bitset[nid >> 3] >> (nid & 7)) & 1


class Pass2Builder(osmium.SimpleHandler):
    """Pass 2: collect matching nodes, then process ways."""
    def __init__(self, bitset_test, outdir):
        super().__init__()
        self.bitset_test = bitset_test
        self.outdir = outdir

        self.node_ids = np.empty(200_000_000, dtype=np.uint64)
        self.node_lons = np.empty(200_000_000, dtype=np.float32)
        self.node_lats = np.empty(200_000_000, dtype=np.float32)
        self.node_count = 0
        self.node_phase = True

        self.sorted_ids = None
        self.sorted_lons = None
        self.sorted_lats = None

        self.points = {}
        self.way_count = 0
        self.point_count = 0
        self.miss_count = 0

    def _ensure_capacity(self):
        if self.node_count < len(self.node_ids):
            return
        new_size = len(self.node_ids) * 2
        sys.stderr.write(f"  Resizing node arrays to {new_size//1_000_000}M...\n")
        self.node_ids = np.resize(self.node_ids, new_size)
        self.node_lons = np.resize(self.node_lons, new_size)
        self.node_lats = np.resize(self.node_lats, new_size)

    def node(self, n):
        if not self.node_phase:
            return
        if self.bitset_test(n.id):
            self._ensure_capacity()
            self.node_ids[self.node_count] = n.id
            self.node_lons[self.node_count] = n.location.lon
            self.node_lats[self.node_count] = n.location.lat
            self.node_count += 1
            if self.node_count % 50_000_000 == 0:
                sys.stderr.write(f"  pass2 nodes: {self.node_count//1_000_000}M\n")

    def _build_index(self):
        sys.stderr.write(f"  Sorting {self.node_count//1_000_000}M nodes...\n")
        t0 = time.time()
        ids = self.node_ids[:self.node_count]
        lons = self.node_lons[:self.node_count]
        lats = self.node_lats[:self.node_count]
        self.node_ids = self.node_lons = self.node_lats = None

        order = np.argsort(ids)
        self.sorted_ids = ids[order]
        self.sorted_lons = lons[order]
        self.sorted_lats = lats[order]
        del ids, lons, lats, order
        sys.stderr.write(f"  Sort done in {time.time()-t0:.1f}s, {len(self.sorted_ids)//1_000_000}M nodes indexed\n")

    def _lookup(self, nid):
        idx = np.searchsorted(self.sorted_ids, nid)
        if idx < len(self.sorted_ids) and self.sorted_ids[idx] == nid:
            return float(self.sorted_lons[idx]), float(self.sorted_lats[idx])
        return None, None

    def way(self, w):
        if self.node_phase:
            self.node_phase = False
            self._build_index()

        if "highway" not in w.tags:
            return
        speed = parse_speed(w.tags.get("maxspeed", ""))
        if speed is None:
            return

        self.way_count += 1
        for n in w.nodes:
            lon, lat = self._lookup(n.ref)
            if lon is None:
                self.miss_count += 1
                continue
            self._add_point(lon, lat, speed)

        if self.way_count % 1_000_000 == 0:
            sys.stderr.write(f"  pass2 ways: {self.way_count//1_000_000}M  points: {self.point_count//1_000_000}M\n")

    def _add_point(self, lon, lat, speed):
        cx = int(math.floor(lon / CELL))
        cy = int(math.floor(lat / CELL))
        lonq = int(round((lon - cx * CELL) / CELL * 65535))
        latq = int(round((lat - cy * CELL) / CELL * 65535))
        lonq = max(0, min(65535, lonq))
        latq = max(0, min(65535, latq))
        sx = min(SUB - 1, lonq * SUB // 65536)
        sy = min(SUB - 1, latq * SUB // 65536)
        key = (cx, cy)
        cell = self.points.get(key)
        if cell is None:
            cell = [None] * (SUB * SUB)
            self.points[key] = cell
        sub = sx + sy * SUB
        buf = cell[sub]
        if buf is None:
            buf = bytearray()
            cell[sub] = buf
        buf += struct.pack("<HHB", lonq, latq, max(0, min(255, speed)))

    def write_grid(self):
        sys.stderr.write(f"\nWriting grid...\n")
        os.makedirs(self.outdir, exist_ok=True)
        index = {}
        total = 0
        for (cx, cy), cell in self.points.items():
            count = sum(0 if b is None else len(b) // 5 for b in cell)
            offsets = [count] * (SUB * SUB)
            acc = 0
            for sub, b in enumerate(cell):
                if b is None:
                    continue
                offsets[sub] = acc
                acc += len(b) // 5
            for s in range(SUB * SUB - 2, -1, -1):
                if offsets[s] == count:
                    offsets[s] = offsets[s + 1]
            buf = bytearray()
            buf += struct.pack("<i", count)
            buf += struct.pack("<%di" % (SUB * SUB), *offsets)
            for b in cell:
                if b is not None:
                    buf += b
            fname = f"{cx}_{cy}.bin.gz"
            path = os.path.join(self.outdir, fname)
            with open(path, "wb") as f:
                f.write(gzip.compress(bytes(buf), 6))
            index[fname] = {"x": cx, "y": cy, "count": count}
            total += count
            if len(index) % 500 == 0:
                sys.stderr.write(f"  cells: {len(index)}  points: {total//1_000_000}M\n")

        with open(os.path.join(self.outdir, "index.json"), "w") as f:
            json.dump(index, f)
        with open(os.path.join(self.outdir, "meta.json"), "w") as f:
            json.dump({"cell": CELL, "sub": SUB, "step_m": 0, "points": total,
                        "files": len(index)}, f)
        sys.stderr.write(f"OK: {total} points in {len(index)} cells -> {self.outdir}\n")


def main():
    if len(sys.argv) != 3:
        print(f"usage: {sys.argv[0]} <planet.pbf> <outdir>")
        sys.exit(1)
    pbf_path = sys.argv[1]
    outdir = sys.argv[2]

    # === PASS 1 ===
    sys.stderr.write("=== PASS 1: Collect node IDs from highway+maxspeed ways ===\n")
    t0 = time.time()
    c = Pass1Collector()
    c.apply_file(pbf_path, locations=False)
    sys.stderr.write(f"Pass 1 done: {c.way_count} ways, {time.time()-t0:.0f}s\n")

    # === PASS 2 ===
    sys.stderr.write("=== PASS 2: Collect nodes + process ways ===\n")
    t1 = time.time()
    b = Pass2Builder(c.test, outdir)
    b.apply_file(pbf_path, locations=False)
    if not b.node_phase:
        pass
    else:
        sys.stderr.write("WARNING: no ways processed (node phase never ended)\n")
    sys.stderr.write(f"Pass 2 done: {b.node_count} nodes, {b.way_count} ways, {b.point_count} points, {time.time()-t1:.0f}s\n")

    # === WRITE GRID ===
    t2 = time.time()
    b.write_grid()
    sys.stderr.write(f"Grid written in {time.time()-t0:.0f}s total\n")

if __name__ == "__main__":
    main()