#!/usr/bin/env python3
"""
Address search (step 2.4): OpenStreetMap places and named streets in South Africa, from OpenStreetMap ONLY. This file
is the published method for the derived database offered under the ODbL at unmapped.co.za/open (ODbL 4.6).

Nothing else is joined in: no Municipal Demarcation Board ward or municipality, no Stats SA code, no other source.
The site works out a result's ward and municipality when someone searches, from our separate MDB ward layer, so the
two databases sit side by side (ODbL 4.5a) and the MDB's no-sale term never enters this one (ODbL 4.4d).

Needs pyosmium and shapely in a virtual env outside the repo:
    python3 -m venv <scratch>/venv && <scratch>/venv/bin/pip install osmium shapely
    <scratch>/venv/bin/python -I scripts/db/address-osm.py --pbf <south-africa-YYMMDD.osm.pbf> --out <dir>

Writes <dir>/osm-places.ndjson, <dir>/osm-streets.ndjson (one JSON object a line) and <dir>/osm-stats.json.

What it keeps (public place and street names only, never personal information; no contributor names, ids or history,
which Geofabrik's extracts don't carry anyway):
  - places: city, town, village, hamlet, suburb, quarter, neighbourhood, borough, locality, as nodes or areas, inside
    South Africa's outline from the same extract (the admin_level 2 boundary tagged ISO3166-1=ZA). Not farms or
    isolated dwellings (often a family's name), and not a lodge, resort or hotel mapped as a village.
  - streets: named public roads (motorway to residential, living and pedestrian streets; a service road only when it
    has no service=* subtype and its name ends in a street word); nothing private or closed, no tracks or paths, no
    names that are numbers or route refs. Same-named ways within ~150 m are one street; a street longer than 1.5 km is
    cut into pieces of about 1 km, each with its middle point, its end points, and the places it runs past.
"""
import argparse
import collections
import json
import math
import re
import sys
import unicodedata

import numpy as np
import osmium
import shapely
import shapely.affinity
from shapely import STRtree, contains_xy, prepare
from shapely.geometry import LineString, Point, shape
import shapely.ops
from shapely.ops import linemerge

PLACE_KINDS = {'city', 'town', 'village', 'hamlet', 'suburb', 'quarter', 'neighbourhood', 'borough', 'locality'}
# Kept as context for streets (the suburb or town a street is in) and as places.
LOCAL_KINDS = PLACE_KINDS
TOWN_KINDS = {'city', 'town'}
ROAD_CLASSES = {
    'motorway': 9, 'trunk': 8, 'primary': 7, 'secondary': 6, 'tertiary': 5, 'unclassified': 4, 'residential': 3,
    'living_street': 3, 'pedestrian': 3, 'road': 2, 'service': 1,
}
CLOSED_ACCESS = {'private', 'no', 'permit', 'customers', 'delivery', 'military', 'agricultural', 'forestry'}
STREET_WORDS = r'(street|road|avenue|drive|crescent|close|lane|way|place|boulevard|highway|terrace|circle|square|walk|loop|mews|row|court|grove|rise|link|parade|promenade|esplanade|straat|weg|laan|singel|rylaan|plein|steeg|hof|pad)'
STREET_END = re.compile(rf'\b{STREET_WORDS}\.?$', re.I)
# A "locality" that is a point of interest rather than a settlement.
POI_LOCALITY = re.compile(r'\b(gate|berth|view ?point|waterhole|water hole|hide|picnic|camp ?site|lookout|look-out|parking|car park|jetty|slipway|quay|wharf|pier|terminal|mine|shaft|dam wall|weir|hotel|lodge|guest ?house|farm|resort|spa|lighthouse|beacon|trig|summit|peak|pass|poort|kloof|bay|beach|point|cave|crossing|junction|interchange|off-?ramp|on-?ramp|circle)\b', re.I)
# A business mapped as a village, hamlet, neighbourhood or locality (a lodge, resort, hotel or campsite): not a
# settlement, and often a private owner's trading name.
BUSINESS = re.compile(r"\b(lodge|eco-lodge|resort|hotel|motel|campsite|camp site|caravan park|chalets?|cottages|guest ?house|guest ?farm|backpackers|b ?& ?b|stables|spa|retreat|holiday)\b", re.I)
SMALL_KINDS = {'village', 'hamlet', 'neighbourhood', 'locality'}
GENERIC = {'informal settlement', 'unnamed', 'unnamed road', 'no name', 'noname', 'road', 'street', 'avenue', 'drive', 'lane', 'private road',
           'access road', 'service road', 'driveway', 'unknown', 'track', 'path', 'worker housing', 'staff housing', 'camp', 'police camp',
           'hostel', 'compound', 'residential', 'village', 'suburb', 'town', 'township', 'settlement', 'farm', 'n/a', 'na', 'none', 'yes', 'tbc', 'test'}
# Names other than `name` that we keep as searchable alternatives: former and other official names, and the name in
# South Africa's languages (English, Afrikaans, isiZulu, isiXhosa, Sesotho, Setswana, Sepedi, siSwati, Tshivenda,
# Xitsonga, isiNdebele).
ALT_KEYS = ['official_name', 'old_name', 'alt_name', 'short_name', 'loc_name', 'name:en', 'name:af', 'name:zu', 'name:xh', 'name:st',
            'name:tn', 'name:nso', 'name:ss', 'name:ve', 'name:ts', 'name:nr', 'old_name:en', 'old_name:af', 'alt_name:en', 'alt_name:af',
            'name:old']
ABBREV = {'st': 'street', 'str': 'street', 'rd': 'road', 'ave': 'avenue', 'av': 'avenue', 'dr': 'drive', 'cres': 'crescent', 'cr': 'crescent',
          'ln': 'lane', 'pl': 'place', 'blvd': 'boulevard', 'hwy': 'highway', 'cl': 'close', 'ext': 'extension', 'tce': 'terrace'}
LAT_KM = 110.574
COS = math.cos(math.radians(29))  # scales longitude to roughly kilometres-equal units across South Africa


def fold(s):
    """A name for grouping: lower case, no accents or apostrophes, words separated by single spaces."""
    s = unicodedata.normalize('NFKD', s)
    s = ''.join(c for c in s if not unicodedata.combining(c))
    s = re.sub(r"['’‘`´ʼ!]", '', s.lower())
    s = re.sub(r'[^a-z0-9]+', ' ', s).strip()
    words = s.split(' ')
    if len(words) > 1 and words[-1] in ABBREV:
        words[-1] = ABBREV[words[-1]]
    return ' '.join(words)


def latin(s):
    return all(ord(c) < 0x250 or not c.isalpha() for c in s)


def ok_name(name):
    if not name or len(name) > 120:
        return False
    f = fold(name)
    if not f or f in GENERIC:
        return False
    if sum(c.isalpha() for c in f) < 2:
        return False
    # Refs and codes: "94006 Track", "R102", "M5", "N2" alone (a numbered route without a name).
    if re.fullmatch(r'[a-z]{0,3} ?\d[\d ]*( track| road)?', f):
        return False
    return latin(name)


def alt_names(tags, name):
    out = []
    seen = {fold(name)}
    for k in ALT_KEYS:
        v = tags.get(k)
        if not v:
            continue
        for part in v.split(';'):
            part = part.strip()
            if part and ok_name(part) and fold(part) not in seen:
                seen.add(fold(part))
                out.append({'name': part, 'type': k})
    return out


def length_m(coords):
    total = 0.0
    for (x1, y1), (x2, y2) in zip(coords, coords[1:]):
        c = math.cos(math.radians((y1 + y2) / 2))
        total += math.hypot((x2 - x1) * 111.320 * c, (y2 - y1) * LAT_KM)
    return total * 1000


def geom_length_m(g):
    if g.geom_type == 'LineString':
        return length_m(list(g.coords))
    if hasattr(g, 'geoms'):
        return sum(geom_length_m(p) for p in g.geoms)
    return 0.0


def lines_of(g):
    if g.is_empty:
        return []
    if g.geom_type == 'LineString':
        return [g]
    if g.geom_type == 'MultiLineString':
        return list(g.geoms)
    if g.geom_type == 'GeometryCollection':
        return [p for x in g.geoms for p in lines_of(x)]
    return []




STEP_M = 1000      # a long street is cut into pieces of about this length, each found by the places it runs past
CLUSTER_DEG = 0.0015   # same-named ways closer than ~150 m are one street


def za_outline(pbf):
    """South Africa's land outline from the extract itself: the admin_level 2 boundary tagged ISO3166-1=ZA."""
    wkb = osmium.geom.WKBFactory()
    fp = (osmium.FileProcessor(pbf).with_locations()
          .with_areas(osmium.filter.TagFilter(('admin_level', '2')))
          .with_filter(osmium.filter.TagFilter(('admin_level', '2'))))
    for o in fp:
        if o.is_area() and o.tags.get('ISO3166-1') == 'ZA' and o.tags.get('boundary') == 'administrative':
            g = shapely.from_wkb(bytes.fromhex(wkb.create_multipolygon(o)))
            return g if g.is_valid else shapely.make_valid(g)
    sys.exit('no South Africa boundary (admin_level 2, ISO3166-1=ZA) in the extract')


def pt(x, y):
    return [round(y, 6), round(x, 6)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--pbf', required=True)
    ap.add_argument('--out', required=True)
    a = ap.parse_args()

    za = za_outline(a.pbf)
    prepare(za)
    inside = lambda x, y: bool(contains_xy(za, x, y))
    print('South Africa outline read', file=sys.stderr)

    stats = collections.Counter()
    nodes, areas, ways = [], [], []
    wkb = osmium.geom.WKBFactory()
    fp = (osmium.FileProcessor(a.pbf)
          .with_locations()
          .with_areas(osmium.filter.KeyFilter('place'))
          .with_filter(osmium.filter.KeyFilter('place', 'highway')))
    for o in fp:
        t = o.tags
        name = t.get('name')
        if o.is_node():
            kind = t.get('place')
            if not kind:
                continue
            if not name:
                stats['place dropped: no name'] += 1
                continue
            if kind not in PLACE_KINDS:
                stats[f'place dropped: {kind}'] += 1
                continue
            if not ok_name(name):
                stats['place dropped: not a usable name'] += 1
                continue
            if kind == 'locality' and POI_LOCALITY.search(name):
                stats['place dropped: locality is a point of interest'] += 1
                continue
            if kind in SMALL_KINDS and BUSINESS.search(name) and 'informal settlement' not in name.lower():
                stats['place dropped: a lodge, resort, hotel or campsite, not a settlement'] += 1
                continue
            nodes.append({'osm': f'n{o.id}', 'kind': kind, 'name': name, 'alt': alt_names(t, name), 'x': o.location.lon, 'y': o.location.lat})
        elif o.is_area():
            kind = t.get('place')
            if not kind or not name or 'highway' in t:
                continue
            if kind not in PLACE_KINDS:
                stats[f'place area dropped: {kind}'] += 1
                continue
            if not ok_name(name) or (kind == 'locality' and POI_LOCALITY.search(name)) or (kind in SMALL_KINDS and BUSINESS.search(name) and 'informal settlement' not in name.lower()):
                stats['place area dropped: name'] += 1
                continue
            try:
                g = shapely.from_wkb(bytes.fromhex(wkb.create_multipolygon(o)))
            except Exception:
                stats['place area dropped: broken geometry'] += 1
                continue
            if not g.is_valid:
                g = shapely.make_valid(g)
            src = f'{"w" if o.from_way() else "r"}{o.orig_id()}'
            areas.append({'osm': src, 'kind': kind, 'name': name, 'alt': alt_names(t, name), 'geom': g})
        elif o.is_way():
            hw = t.get('highway')
            if not hw:
                continue
            if not name:
                continue
            if hw not in ROAD_CLASSES:
                stats[f'street dropped: highway={hw}'] += 1
                continue
            if t.get('access') in CLOSED_ACCESS or t.get('motor_vehicle') in ('private', 'no') and hw == 'service':
                stats['street dropped: private or closed access'] += 1
                continue
            if hw == 'service' and (t.get('service') or not STREET_END.search(name or '')):
                stats['street dropped: service road (driveway, parking, or not named as a street)'] += 1
                continue
            if t.get('area') == 'yes':
                stats['street dropped: an area, not a way'] += 1
                continue
            if not ok_name(name):
                stats['street dropped: not a usable name'] += 1
                continue
            try:
                coords = [(n.lon, n.lat) for n in o.nodes]
            except osmium.InvalidLocationError:
                stats['street dropped: nodes missing from the extract'] += 1
                continue
            if len(coords) < 2:
                continue
            notable = bool(t.get('wikidata') or t.get('wikipedia'))
            ways.append((o.id, name, hw, alt_names(t, name), coords, notable))
    print(f'read: {len(nodes)} place nodes, {len(areas)} place areas, {len(ways)} street ways', file=sys.stderr)

    # ── Places ─────────────────────────────────────────────────────────────────────────────────────────────
    # An area and a node of the same name (the node inside the area or within 1 km) are one place: the node's point.
    node_tree = STRtree([Point(n['x'], n['y']) for n in nodes]) if nodes else None
    used = set()
    places = []
    for ar in areas:
        g = ar['geom']
        match = None
        if node_tree is not None:
            for i in node_tree.query(g, predicate='dwithin', distance=0.01):
                n = nodes[int(i)]
                if fold(n['name']) == fold(ar['name']) and int(i) not in used:
                    match = int(i)
                    break
        if match is not None:
            used.add(match)
            n = nodes[match]
            n['area'] = g
            n['osm_area'] = ar['osm']
            n['alt'] += [x for x in ar['alt'] if fold(x['name']) not in {fold(y['name']) for y in n['alt']}]
            stats['place areas merged with their node'] += 1
        else:
            p = g.representative_point()
            places.append({**{k: v for k, v in ar.items() if k != 'geom'}, 'x': p.x, 'y': p.y, 'area': g})
    places = nodes + places

    order = ['city', 'town', 'borough', 'suburb', 'village', 'quarter', 'neighbourhood', 'hamlet', 'locality']
    places.sort(key=lambda p: (order.index(p['kind']), -len(p['alt']), 0 if 'area' in p else 1))
    ptree = STRtree([Point(p['x'] * COS, p['y']) for p in places])
    dropped = set()
    for i, p in enumerate(places):
        if i in dropped:
            continue
        for j in ptree.query(Point(p['x'] * COS, p['y']), predicate='dwithin', distance=2 / 111):
            j = int(j)
            if j != i and j not in dropped and j > i and fold(places[j]['name']) == fold(p['name']):
                dropped.add(j)
    stats['place duplicates merged (same name within 2 km)'] = len(dropped)
    places = [p for i, p in enumerate(places) if i not in dropped]

    out_places = []
    for p in places:
        if not inside(p['x'], p['y']):
            stats['place dropped: outside South Africa (OSM national boundary)'] += 1
            continue
        rec = {'id': f"p-{p['osm']}", 'osm': p['osm'], 'kind': p['kind'], 'name': p['name'], 'alt': p['alt'], 'at': pt(p['x'], p['y'])}
        if 'osm_area' in p:
            rec['osm_area'] = p['osm_area']
        if 'area' in p:
            g = p['area']
            # Its extent: the area's points furthest north, south, east and west (where a big place reaches).
            xs, ys = [], []
            for poly in (g.geoms if hasattr(g, 'geoms') else [g]):
                if poly.geom_type != 'Polygon':
                    continue
                c = list(poly.exterior.coords)
                xs += [q[0] for q in c]
                ys += [q[1] for q in c]
            if xs:
                ext = [(min(range(len(xs)), key=lambda i: f(i))) for f in (lambda i: ys[i], lambda i: -ys[i], lambda i: xs[i], lambda i: -xs[i])]
                rec['reach'] = [pt(xs[i], ys[i]) for i in ext]
            rec['area_km2'] = round(g.area * 111.32 * LAT_KM * math.cos(math.radians(p['y'])), 2)
        out_places.append(rec)
    print(f'places: {len(out_places)}', file=sys.stderr)

    def tree_of(kinds):
        pts = [p for p in out_places if p['kind'] in kinds]
        return pts, STRtree([Point(p['at'][1] * COS, p['at'][0]) for p in pts])

    local_pts, local_tree = tree_of(LOCAL_KINDS)
    town_pts, town_tree = tree_of(TOWN_KINDS)
    vil_pts, vil_tree = tree_of({'village'})

    def nearest(pts, tree, x, y, km, exclude=None):
        best, bd = None, 1e9
        for i in tree.query(Point(x * COS, y), predicate='dwithin', distance=km / 111):
            p = pts[int(i)]
            if exclude and p['osm'] == exclude:
                continue
            d = (p['at'][1] * COS - x * COS) ** 2 + (p['at'][0] - y) ** 2
            if d < bd:
                best, bd = p, d
        return best

    def town_of(x, y, exclude=None):
        t = nearest(town_pts, town_tree, x, y, 20, exclude) or nearest(vil_pts, vil_tree, x, y, 6, exclude)
        return t['name'] if t else None

    for p in out_places:
        if p['kind'] not in TOWN_KINDS:
            p['town'] = town_of(p['at'][1], p['at'][0], p['osm'])

    # ── Streets ────────────────────────────────────────────────────────────────────────────────────────────
    # Same-named ways within ~150 m of each other are one street. A street longer than STEP_M is cut into pieces of
    # about STEP_M, so each piece is found by the suburb it runs through and has a point near the place asked about.
    by_name = collections.defaultdict(list)
    for w in ways:
        line = LineString(w[4])
        if line.length > 0:
            by_name[fold(w[1])].append((w, line))
    out_streets = []
    for key, items in by_name.items():
        geoms = [l for _, l in items]
        parent = list(range(len(geoms)))

        def find(i):
            while parent[i] != i:
                parent[i] = parent[parent[i]]
                i = parent[i]
            return i
        if len(geoms) > 1:
            t = STRtree(geoms)
            a_idx, b_idx = t.query(geoms, predicate='dwithin', distance=CLUSTER_DEG)
            for x, y in zip(a_idx, b_idx):
                rx, ry = find(int(x)), find(int(y))
                if rx != ry:
                    parent[rx] = ry
        clusters = collections.defaultdict(list)
        for i in range(len(geoms)):
            clusters[find(i)].append(i)
        for members in clusters.values():
            ms = [items[i][0] for i in members]
            names = collections.Counter(m[1] for m in ms)
            name = max(names.items(), key=lambda kv: (kv[1], len(kv[0])))[0]
            hw = max((m[2] for m in ms), key=lambda h: ROAD_CLASSES[h])
            alt, seen = [], {fold(name)}
            for m in ms:
                for x in m[3] + [{'name': n, 'type': 'name'} for n in names]:
                    if fold(x['name']) not in seen:
                        seen.add(fold(x['name']))
                        alt.append(x)
            merged = linemerge([items[i][1] for i in members])
            comps = sorted(lines_of(merged), key=lambda g: -g.length)
            total = sum(geom_length_m(c) for c in comps)
            if total < 5:
                stats['streets dropped: under 5 m'] += 1
                continue
            first_way = min(m[0] for m in ms)
            notable = any(m[5] for m in ms)
            # The pieces: the whole street when short; else each part cut to about STEP_M.
            pieces = []
            if total <= STEP_M * 1.5:
                pieces.append(comps)
            else:
                for c in comps:
                    lm = geom_length_m(c)
                    n = max(1, round(lm / STEP_M))
                    for k in range(n):
                        pieces.append([shapely.ops.substring(c, k / n, (k + 1) / n, normalized=True)])
            pieces = [pc for pc in pieces if sum(geom_length_m(g) for g in pc) >= 5]
            for n_piece, pc in enumerate(pieces):
                main_line = max(pc, key=lambda g: g.length)
                mid = main_line.interpolate(0.5, normalized=True)
                if not inside(mid.x, mid.y):
                    stats['street pieces dropped: outside South Africa'] += 1
                    continue
                ends = []
                for g in pc[:3]:
                    for frac in (0, 1):
                        q = g.interpolate(frac, normalized=True)
                        if inside(q.x, q.y):
                            ends.append(pt(q.x, q.y))
                line = shapely.affinity.scale(shapely.geometry.MultiLineString([list(g.coords) for g in pc]), xfact=COS, yfact=1, origin=(0, 0))
                loc = nearest(local_pts, local_tree, mid.x, mid.y, 3)
                along = [local_pts[int(i)] for i in local_tree.query(line, predicate='dwithin', distance=1 / 111)]
                along.sort(key=lambda p: line.distance(Point(p['at'][1] * COS, p['at'][0])))
                near = [p['name'] for p in along if p is not loc][:4]
                rec = {'id': f's-w{first_way}' + (f'-{n_piece + 1}' if len(pieces) > 1 else ''), 'kind': 'street', 'name': name, 'alt': alt, 'class': hw,
                       'at': pt(mid.x, mid.y), 'ends': ends, 'length_m': round(sum(geom_length_m(g) for g in pc)),
                       'locality': loc['name'] if loc else None, 'near': near, 'town': town_of(mid.x, mid.y),
                       'ways': sorted({m[0] for m in ms})[:50], 'notable': notable}
                out_streets.append(rec)
    print(f'streets: {len(out_streets)}', file=sys.stderr)

    with open(f'{a.out}/osm-places.ndjson', 'w', encoding='utf8') as f:
        for p in sorted(out_places, key=lambda p: p['id']):
            f.write(json.dumps(p, ensure_ascii=False) + '\n')
    with open(f'{a.out}/osm-streets.ndjson', 'w', encoding='utf8') as f:
        for s in sorted(out_streets, key=lambda s: s['id']):
            f.write(json.dumps(s, ensure_ascii=False) + '\n')
    json.dump({'counts': dict(stats)}, open(f'{a.out}/osm-stats.json', 'w'), indent=1, ensure_ascii=False)
    print(json.dumps(dict(stats), indent=1, ensure_ascii=False), file=sys.stderr)


if __name__ == '__main__':
    main()