#!/usr/bin/env python3
"""Construit une grille binaire des maxspeed OSM (par route, échantillonnée en points).

Pipeline :
  osmium tags-filter <extrait>.osm.pbf w/highway -o <extrait>-roads.osm.pbf
  python3 build_maxspeed.py <extrait>-roads.osm.pbf <sortie> [--cache=<file>]

Sortie : <sortie>/grid.dat + <sortie>/index.json + <sortie>/meta.json
  - Chaque cellule = 0.5° x 0.5° (lat, lon).
  - Blob par cellule (gzip) : int32 count ; 32x32 sous-cellules => int32 offset[1024] ;
    puis count points de (lonq:int16, latq:int16, speed:int8) rangés par sous-cellule.
  - lonq/latq : quantifiés sur 0..65535 dans la cellule (precision ~0.8 m).
"""
import sys, os, json, gzip, math, struct, re
import osmium
import sqlite3

CELL = 0.5          # degrés par cellule
SUB = 32            # sous-cellules par dimension (grille 32x32)
STEP_M_DEFAULT = 40.0       # pas d'échantillonnage le long des routes (m)
STEP_M = STEP_M_DEFAULT

# maxspeed textuel -> km/h (valeurs par défaut OSM)
ZONE = {
    "none": None,
    "signals": None,
    "variable": None,
    "walk": 5,
    "living_street": 20,
    "urban": 50, "rural": 80, "trunk": 110, "motorway": 130,
}
# codes pays -> (urban, rural) km/h ; motorway sans maxspeed = None (ex. DE)
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),
}
# Les highways de service/privées sans maxspeed sont exclues : on ne garde que les
# ways marqués maxspeed (ou avec une valeur par défaut de zone reconnue).

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
    # code pays type FR:rural / DE:urban
    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

class MaxSpeedHandler(osmium.SimpleHandler):
    def __init__(self):
        super().__init__()
        self.points = {}   # (cellx, celly) -> dict subcell -> list of (lonq, latq, speed)
        self.skipped = 0

    def way(self, w):
        if "highway" not in w.tags:
            return
        speed = parse_speed(w.tags.get("maxspeed", ""))
        if speed is None:
            return
        locs = []
        for n in w.nodes:
            loc = n.location
            if loc is not None:
                locs.append((loc.lon, loc.lat))
        if len(locs) < 2:
            return
        self.process_locs(locs, speed)

    def process_locs(self, locs, speed):
        if STEP_M <= 0:
            for (lon, lat) in locs:
                self.add_point(lon, lat, speed)
            return
        prev = locs[0]
        lon, lat = prev
        acc = 0.0
        emit = [(lon, lat)]
        for node in locs[1:]:
            a = haversine(prev[1], prev[0], node[1], node[0])
            acc += a
            prev = node
            while acc >= STEP_M:
                f = 1.0 - (acc - STEP_M) / a if a > 0 else 0.0
                ilon = lon + (node[0] - lon) * f
                ilat = lat + (node[1] - lat) * f
                emit.append((ilon, ilat))
                acc -= STEP_M
            lon, lat = node[0], node[1]
        if (lon, lat) != emit[-1]:
            emit.append((lon, lat))
        for (el, ea) in emit:
            self.add_point(el, ea, speed)

    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 haversine(lat1, lon1, lat2, lon2):
    R = 6371000.0
    p1 = math.radians(lat1); p2 = math.radians(lat2)
    dp = math.radians(lat2 - lat1); dl = math.radians(lon2 - lon1)
    a = math.sin(dp / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dl / 2) ** 2
    return R * 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))

GP_FLAGS = {"byteorder": 0x01, "envelope": 0x02, "empty": 0x10}

def parse_gp_wkb(raw):
    """GeoPackage binary -> liste (lon, lat) pour LineString (ou Multi)."""
    if not isinstance(raw, (bytes, bytearray)) or len(raw) < 8 or raw[:2] != b"GP":
        return None
    flags = raw[3]
    off = 8
    if flags & GP_FLAGS["envelope"]:
        etype = (flags >> 2) & 0x03
        if etype != 0:
            return None
        off += 4 * 8
    bo = "<" if flags & GP_FLAGS["byteorder"] else ">"
    if off + 5 > len(raw):
        return None
    gtype = struct.unpack(bo + "i", raw[off + 1:off + 5])[0]
    if gtype == 2:   # LineString
        return parse_wkb_points(raw, off, bo)
    if gtype == 5:   # MultiLineString
        out = []
        num = struct.unpack(bo + "i", raw[off + 5:off + 9])[0]
        p = off + 9
        for _ in range(num):
            pts = parse_wkb_points(raw, p, bo)
            if pts:
                out.extend(pts)
            p += 4 + 4 * 8 * (len(pts) if pts else 0)
        return out if len(out) >= 2 else None
    return None

def parse_wkb_points(raw, off, bo):
    n = struct.unpack(bo + "i", raw[off + 5:off + 9])[0]
    out = []
    p = off + 9
    for _ in range(n):
        out.append((struct.unpack(bo + "d", raw[p:p + 8])[0],
                    struct.unpack(bo + "d", raw[p + 8:p + 16])[0]))
        p += 16
    return out

MS_RE = re.compile(r'"maxspeed"=>"([^"]*)"')

def parse_other_tags(s):
    m = MS_RE.search(s or "")
    return m.group(1) if m else None

def load_gpkg(h, gpkg):
    con = sqlite3.connect(gpkg)
    cur = con.execute("SELECT geom, highway, other_tags FROM lines")
    n = 0
    for raw, hw, tags in cur:
        if not hw or "maxspeed" not in (tags or ""):
            continue
        v = parse_other_tags(tags)
        speed = parse_speed(v)
        if speed is None:
            continue
        locs = parse_gp_wkb(raw)
        if not locs or len(locs) < 2:
            continue
        h.process_locs(locs, speed)
        n += 1
        if n % 500000 == 0:
            sys.stderr.write("ways maxspeed: %dM\n" % (n // 1000000))
    con.close()
    return n

def build(src_pbf, out_dir, cache_file=None):
    os.makedirs(out_dir, exist_ok=True)
    h = MaxSpeedHandler()
    if src_pbf.endswith(".gpkg") or src_pbf.endswith(".db"):
        n = load_gpkg(h, src_pbf)
        print(f"GPKG : {n} ways traités")
    else:
        if cache_file:
            idx = osmium.index.create_map("sparse_file_array,%s" % cache_file)
            h.apply_file(src_pbf, locations=osmium.NodeLocationsForWays(idx))
        else:
            h.apply_file(src_pbf, locations=True)
    # écriture : un fichier par cellule (nom: X_Y.bin.gz) + index
    index = {}
    total = 0
    for (cx, cy), cell in h.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
        # remplissage arrière : les sous-cellules vides pointent vers la suivante
        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"
        with open(os.path.join(out_dir, fname), "wb") as f:
            f.write(gzip.compress(bytes(buf), 6))
        index[fname] = {"x": cx, "y": cy, "count": count}
        total += count
    with open(os.path.join(out_dir, "index.json"), "w") as f:
        json.dump(index, f)
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump({"cell": CELL, "sub": SUB, "step_m": STEP_M, "points": total,
                   "files": len(index)}, f)
    print(f"OK : {total} points dans {len(index)} cellules -> {out_dir}")

if __name__ == "__main__":
    cache = None
    args = sys.argv[1:]
    extra = {}
    while args and args[-1].startswith("--"):
        a = args.pop()
        if a.startswith("--cache="):
            cache = a.split("=", 1)[1]
        elif a == "--no-sample":
            STEP_M = 0
        else:
            print(f"Option inconnue: {a}")
            sys.exit(1)
    if len(args) != 2:
        print("usage: build_maxspeed.py <extrait.osm.pbf|gpkg> <out_dir> [--cache=<file>] [--no-sample]")
        sys.exit(1)
    build(args[0], args[1], cache)