"""Build the area road graphs (JSON) for the prototype from a Geofabrik .osm.pbf.

Usage:
  python -I tools/extract_graph.py <greater-london.osm.pbf> --all                    every region in tools/regions.json
  python -I tools/extract_graph.py <greater-london.osm.pbf> --region=east[,west…]    some regions
      → osm/<id>.json + osm/<id>-pois.json in $TOPO_OSM_DIR, default topo-core/osm/ (bbox: the region's graphBbox, else its bbox; one pass over the pbf).
        hackney also refreshes the old names osm/area.json + osm/pois.json (aliases, same content).
  python -I tools/extract_graph.py <greater-london.osm.pbf> <minlon,minlat,maxlon,maxlat> <out.json> [<pois.json>]
Tiles of a region (same bbox; free Protomaps daily build, no sign-up; z0–10 dropped, the map never goes below z11):
  pmtiles extract https://build.protomaps.com/20261102.pmtiles tiles/<id>.pmtiles --bbox=<w,s,e,n> --minzoom=11 --maxzoom=15
  (hackney.pmtiles: the review-6 extract, z0–15.) Budget: ≤ 12 MB per archive, ≤ 80 MB in all (BOARD D1); use
  --maxzoom=14 for a region that doesn't fit.

Adapted from the TECH prototype's task_graph.py. Changes:
- keeps highway=pedestrian ways (flag PED) so the grader can say "Not a road for cars: {street}";
- keeps `ref` (road number) apart from `name`;
- dual_carriageway flag; service flag; closures are NOT baked in (the grader closes every way
  with the task's closed name at load time, so one area graph serves every task);
- fixes the `except=motorcar` restriction filter; keeps via-node restrictions only (via-way: known gap);
- optional: named POIs (hospitals, stations, parks, museums, schools…) for MCQ "places" items;
- review 7 (M2, M3): tunnel / bridge flags; busonly (OSM bus/taxi-only ways + tools/overrides.json) and park
  (royal park roads, tools/overrides.json); busonly and park edges also get noaccess.
Edge = [a, b, flags, nameIdx, refIdx, length_m, [mid lon,lat,…]]; barriers = node indices that block cars (v3)
"""
import gzip, json, math, os, shutil, sys
import osmium

DRIVABLE = {
    'motorway', 'trunk', 'primary', 'secondary', 'tertiary', 'unclassified', 'residential',
    'motorway_link', 'trunk_link', 'primary_link', 'secondary_link', 'tertiary_link', 'living_street', 'service',
}
PEDESTRIAN = {'pedestrian'}
F = dict(oneway=1, reverse=2, noaccess=4, conditional=8, major=32, layered=64, ped=128, dual=256, service=512, roundabout=1024,
         tunnel=2048, bridge=4096, busonly=8192, park=16384)

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OSM = os.environ.get('TOPO_OSM_DIR') or os.path.join(ROOT, 'osm')  # topo-core/osm/ (git-ignored) unless TOPO_OSM_DIR is set
pbf = sys.argv[1]
JOBS = []  # (region id, bbox, graph out, pois out)
if len(sys.argv) > 2 and sys.argv[2].startswith('--'):
    regions = json.load(open(os.path.join(ROOT, 'tools', 'regions.json')))
    opt = sys.argv[2]
    ids = list(regions) if opt == '--all' else opt.split('=', 1)[1].split(',') if opt.startswith('--region=') else None
    if not ids or any(i not in regions for i in ids):
        sys.exit(f'usage: extract_graph.py <pbf> --all | --region=<id>[,<id>…]  (regions: {", ".join(regions)})')
    os.makedirs(OSM, exist_ok=True)
    for i in ids:
        r = regions[i]
        JOBS.append((i, r.get('graphBbox') or r['bbox'], os.path.join(OSM, f'{i}.json'), os.path.join(OSM, f'{i}-pois.json')))
else:
    JOBS.append((None, list(map(float, sys.argv[2].split(','))), sys.argv[3], sys.argv[4] if len(sys.argv) > 4 else None))
pois_out = any(j[3] for j in JOBS)
U = [min(j[1][0] for j in JOBS), min(j[1][1] for j in JOBS), max(j[1][2] for j in JOBS), max(j[1][3] for j in JOBS)]
OVR = json.load(open(os.path.join(ROOT, 'tools', 'overrides.json')))['rules']
PARKS = {n for r in OVR for n in r.get('within', [])}


def in_box(b, lon, lat):
    return b[0] <= lon <= b[2] and b[1] <= lat <= b[3]


def inside(lon, lat):
    return in_box(U, lon, lat)


def hav(a, b):
    R = 6371008.8
    la1, la2 = math.radians(a[1]), math.radians(b[1])
    dla, dlo = la2 - la1, math.radians(b[0] - a[0])
    h = math.sin(dla / 2) ** 2 + math.cos(la1) * math.cos(la2) * math.sin(dlo / 2) ** 2
    return 2 * R * math.asin(math.sqrt(h))


def in_ring(pt, ring):
    x, y = pt
    c = False
    for i in range(len(ring)):
        x1, y1 = ring[i]
        x2, y2 = ring[i - 1]
        if (y1 > y) != (y2 > y) and x < (x2 - x1) * (y - y1) / (y2 - y1) + x1:
            c = not c
    return c


POI_KEYS = {
    ('amenity', 'hospital'): 'hospital', ('railway', 'station'): 'station', ('leisure', 'park'): 'park',
    ('tourism', 'museum'): 'museum', ('amenity', 'theatre'): 'theatre', ('amenity', 'townhall'): 'town hall',
    ('amenity', 'library'): 'library', ('amenity', 'police'): 'police station', ('amenity', 'fire_station'): 'fire station',
    ('leisure', 'sports_centre'): 'sports centre', ('amenity', 'place_of_worship'): 'place of worship',
    ('amenity', 'school'): 'school', ('amenity', 'college'): 'college', ('tourism', 'hotel'): 'hotel',
}


class H(osmium.SimpleHandler):
    def __init__(self):
        super().__init__()
        self.ways, self.restr, self.pois = [], [], []
        self.barriers = {}
        self.parks = {}

    def poi(self, kind, name, coords, wikidata):
        if coords:
            lon = sum(c[0] for c in coords) / len(coords)
            lat = sum(c[1] for c in coords) / len(coords)
            if inside(lon, lat):
                self.pois.append((lon, lat, dict(kind=kind, name=name, lon=round(lon, 5), lat=round(lat, 5), wd=wikidata)))

    def node(self, n):
        t = n.tags
        b = t.get('barrier')
        if b:
            # barriers that stop cars (QA D18), following the OSRM car profile: a barrier blocks unless it is a
            # passable kind (gate, lift gate…) without a "no" access tag, or it explicitly lets motor vehicles through
            acc = t.get('motor_vehicle') or t.get('motorcar') or t.get('vehicle') or t.get('access')
            passable = b in ('cattle_grid', 'border_control', 'toll_booth', 'sally_port', 'gate', 'lift_gate', 'no', 'entrance',
                             'height_restrictor', 'arch', 'kerb', 'yes')
            if acc in ('yes', 'permissive', 'designated', 'destination', 'customers'):
                passable = True
            if acc in ('no', 'private', 'agricultural', 'forestry', 'delivery', 'emergency'):
                passable = False
            if not passable and inside(n.location.lon, n.location.lat):
                self.barriers[n.id] = (n.location.lon, n.location.lat)
        if not pois_out:
            return
        for (k, v), kind in POI_KEYS.items():
            if t.get(k) == v and t.get('name'):
                self.poi(kind, t.get('name'), [(n.location.lon, n.location.lat)], t.get('wikidata'))

    def way(self, w):
        t = w.tags
        if pois_out:
            for (k, v), kind in POI_KEYS.items():
                if t.get(k) == v and t.get('name'):
                    try:
                        self.poi(kind, t.get('name'), [(n.lon, n.lat) for n in w.nodes], t.get('wikidata'))
                    except osmium.InvalidLocationError:
                        pass
        hw = t.get('highway')
        if hw not in DRIVABLE and hw not in PEDESTRIAN:
            return
        if t.get('area') == 'yes':
            return
        if hw == 'service' and t.get('service') in ('parking_aisle', 'driveway', 'drive-through', 'emergency_access'):
            return
        try:
            coords = [(n.lon, n.lat) for n in w.nodes]
        except osmium.InvalidLocationError:
            return
        if not any(inside(*c) for c in coords):
            return
        acc = t.get('motor_vehicle') or t.get('motorcar') or t.get('vehicle') or t.get('access')
        cond = any((t.get(k) or '').strip().startswith('no') for k in
                   ('motor_vehicle:conditional', 'motorcar:conditional', 'access:conditional', 'vehicle:conditional'))
        ow = t.get('oneway')
        rb = t.get('junction') in ('roundabout', 'circular')
        if ow is None and (rb or hw in ('motorway', 'motorway_link')):
            ow = 'yes'
        # one-way for cars only (QA2 N6): oneway:motorcar|motor_vehicle|vehicle, or *:backward / *:forward = no
        # (a bus-only contraflow is one-way for cars). A car-specific oneway wins over the plain tag.
        for k in ('oneway:motorcar', 'oneway:motor_vehicle', 'oneway:vehicle'):
            if t.get(k) in ('yes', '1', 'true', '-1'):
                ow = t.get(k)
                break
        if ow in (None, 'no'):
            back = any(t.get(f'{k}:backward') == 'no' for k in ('motorcar', 'motor_vehicle', 'vehicle', 'access'))
            fwd = any(t.get(f'{k}:forward') == 'no' for k in ('motorcar', 'motor_vehicle', 'vehicle', 'access'))
            if back and not fwd:
                ow = 'yes'
            elif fwd and not back:
                ow = '-1'
        lay = t.get('layer', '0') or '0'
        self.ways.append(dict(
            id=w.id, nodes=[n.ref for n in w.nodes], coords=coords, hw=hw,
            name=t.get('name') or '', ref=(t.get('ref') or '').split(';')[0], oneway=ow, rb=rb,
            noaccess=acc in ('no', 'private', 'agricultural', 'forestry', 'delivery'),
            cond=cond, dual=t.get('dual_carriageway') == 'yes',
            # bus/taxi-only (a bus gate or a bus street): cars barred, buses, coaches or taxis allowed
            bus=acc == 'no' and any(t.get(k) in ('yes', 'designated') for k in ('bus', 'psv', 'taxi')),
            # levels (M2): a building passage is not a tunnel on the map
            tunnel=t.get('tunnel') not in (None, 'no', 'building_passage'), bridge=t.get('bridge') not in (None, 'no'),
            layer=int(lay) if lay.lstrip('-').isdigit() else 0))

    def area(self, a):
        # royal park outlines for the `within` override rules
        if PARKS and a.tags.get('leisure') == 'park' and a.tags.get('name') in PARKS:
            for outer in a.outer_rings():
                self.parks.setdefault(a.tags.get('name'), []).append([(p.lon, p.lat) for p in outer])

    def relation(self, r):
        t = r.tags
        if t.get('type') != 'restriction':
            return
        kind = t.get('restriction') or t.get('restriction:motorcar') or t.get('restriction:motor_vehicle')
        if not kind or 'motorcar' in (t.get('except') or ''):
            return
        frm = [m.ref for m in r.members if m.role == 'from' and m.type == 'w']
        via = [m.ref for m in r.members if m.role == 'via' and m.type == 'n']
        to = [m.ref for m in r.members if m.role == 'to' and m.type == 'w']
        if frm and via and to:
            self.restr.append((frm[0], via[0], to[0], kind))


h = H()
h.apply_file(pbf, locations=True)
for n in PARKS - set(h.parks):
    print(f'warning: park outline "{n}" not found in the pbf (overrides.json `within`)')


def along(cs, f=0.5):
    """the point at fraction f of a polyline's length"""
    d = [hav(cs[k], cs[k + 1]) for k in range(len(cs) - 1)]
    want, at = sum(d) * f, 0.0
    for k, x in enumerate(d):
        if at + x >= want and x > 0:
            t = (want - at) / x
            return (cs[k][0] + t * (cs[k + 1][0] - cs[k][0]), cs[k][1] + t * (cs[k + 1][1] - cs[k][1]))
        at += x
    return cs[len(cs) // 2]


def override_bits(name, cs):
    """flags our map overrides add to an edge (tools/overrides.json), and the ids of the rules that matched"""
    bits, clear, hit = 0, 0, []
    mid = along(cs)
    for r in OVR:
        if 'clear' in r:  # our map shows this way open although OSM restricts it (e.g. a timed closure tag)
            if name in r['names'] and in_box(r['bbox'], *mid):
                for k in r['clear']:
                    clear |= F[k]
                hit.append(r['id'])
            continue
        b = F[r['set']]
        if 'within' in r:
            ok = any(in_box(rb, *mid) and in_ring(mid, ring) for n in r['within'] for ring, rb in PARK_RINGS.get(n, []))
        else:
            ok = name in r['names'] and in_box(r['bbox'], *mid)
        if ok:
            bits |= b | F['noaccess']
            hit.append(r['id'])
    return bits, clear, hit


PARK_RINGS = {n: [(ring, [min(p[0] for p in ring), min(p[1] for p in ring), max(p[0] for p in ring), max(p[1] for p in ring)]) for ring in rings]
              for n, rings in h.parks.items()}


def build(region, bbox, out, pois_out):
    ways = [w for w in h.ways if any(in_box(bbox, *c) for c in w['coords'])]
    barriers_osm = {n for n, (lon, lat) in h.barriers.items() if in_box(bbox, lon, lat)}
    use = {}
    for w in ways:
        for n in w['nodes']:
            use[n] = use.get(n, 0) + 1
    for n in barriers_osm:  # a barrier is a graph node, so the engine can block passing through it
        if n in use:
            use[n] += 1
    node_ix, nodes = {}, []

    def nid(osm_id, c):
        if osm_id not in node_ix:
            node_ix[osm_id] = len(nodes)
            nodes.append([round(c[0], 6), round(c[1], 6)])
        return node_ix[osm_id]

    names, name_ix = [''], {'': 0}

    def name_id(s):
        if s not in name_ix:
            name_ix[s] = len(names)
            names.append(s)
        return name_ix[s]

    edges, way_edges, rule_hits = [], {}, {}
    for w in ways:
        flags = 0
        if w['oneway'] in ('yes', '1', 'true'): flags |= F['oneway']
        if w['oneway'] == '-1': flags |= F['reverse']
        if w['noaccess'] or w['bus']: flags |= F['noaccess']
        if w['cond']: flags |= F['conditional']
        if w['hw'] in ('motorway', 'trunk', 'primary', 'secondary', 'trunk_link', 'primary_link', 'secondary_link'): flags |= F['major']
        if w['layer'] != 0: flags |= F['layered']
        if w['hw'] in PEDESTRIAN: flags |= F['ped']
        if w['dual']: flags |= F['dual']
        if w['hw'] == 'service': flags |= F['service']
        if w['rb']: flags |= F['roundabout']
        if w['tunnel']: flags |= F['tunnel']
        if w['bridge']: flags |= F['bridge']
        if w['bus']: flags |= F['busonly']
        seg_start = 0
        for i in range(1, len(w['nodes'])):
            if use[w['nodes'][i]] > 1 or i == len(w['nodes']) - 1:
                cs = w['coords'][seg_start:i + 1]
                length = sum(hav(cs[k], cs[k + 1]) for k in range(len(cs) - 1))
                if length < 0.05:
                    seg_start = i
                    continue
                a, b = nid(w['nodes'][seg_start], cs[0]), nid(w['nodes'][i], cs[-1])
                mid = [round(v, 6) for c in cs[1:-1] for v in c]
                ob, oc, hit = override_bits(w['name'], cs)
                for r in hit:
                    rule_hits[r] = rule_hits.get(r, 0) + 1
                way_edges.setdefault(w['id'], []).append(len(edges))
                edges.append([a, b, (flags & ~oc) | ob, name_id(w['name']), name_id(w['ref']), round(length, 1), mid])
                seg_start = i

    restr = []
    for frm, via, to, kind in h.restr:
        if via not in node_ix: continue
        v = node_ix[via]
        fe = [e for e in way_edges.get(frm, []) if v in (edges[e][0], edges[e][1])]
        te = [e for e in way_edges.get(to, []) if v in (edges[e][0], edges[e][1])]
        if fe and te:
            restr.append([fe[0], v, te[0], 1 if kind.startswith('only_') else 0, kind])

    barriers = sorted(node_ix[n] for n in barriers_osm if n in node_ix)
    g = dict(v=3, bbox=list(bbox), names=names, nodes=nodes, edges=edges, restr=restr, barriers=barriers, flags=F,
             snapshot='2026-11-02', attribution='© OpenStreetMap contributors, ODbL 1.0')
    s = json.dumps(g, separators=(',', ':'))
    open(out, 'w').write(s)
    cnt = {k: sum(1 for e in edges if e[2] & F[k]) for k in ('tunnel', 'bridge', 'busonly', 'park')}
    print(f'{out}: {len(nodes)} nodes, {len(edges)} edges, {len(restr)} turn restrictions, {len(barriers)} barriers, '
          f'{len(s)/1024:.0f} KB raw, {len(gzip.compress(s.encode()))/1024:.0f} KB gz; '
          + ', '.join(f'{k} {v}' for k, v in cnt.items()) + (f'; overrides {rule_hits}' if rule_hits else ''))
    for r in OVR:  # a rule whose box lies in this region but matched nothing: a renamed street or a wrong box
        rb = r.get('bbox')
        if rb and rb[0] >= bbox[0] and rb[1] >= bbox[1] and rb[2] <= bbox[2] and rb[3] <= bbox[3] and r['id'] not in rule_hits:
            print(f'warning: override "{r["id"]}" matched no edge in {out}')
    if pois_out:
        seen, pois = set(), []
        for lon, lat, p in h.pois:
            if not in_box(bbox, lon, lat): continue
            k = (p['kind'], p['name'])
            if k in seen: continue
            seen.add(k); pois.append(p)
        json.dump(pois, open(pois_out, 'w'), separators=(',', ':'))
        print(f'{pois_out}: {len(pois)} POIs')
    if region == 'hackney':  # the old names (osm/area.json, osm/pois.json) stay as aliases
        shutil.copyfile(out, os.path.join(os.path.dirname(out), 'area.json'))
        if pois_out:
            shutil.copyfile(pois_out, os.path.join(os.path.dirname(out), 'pois.json'))


for job in JOBS:
    build(*job)