#!/usr/bin/env python3
"""Builds assets/ from the real datasets in data/raw/.

  assets/mask.png       RGB, W x H.  R = land coverage (0-255), G = ferry corridor, B = 0
  assets/cities.json    [[name, lat, lon, population], ...] sorted by population, desc
  assets/overlays.json  {key: {"lines": [[[lon, lat], ...], ...], "km": total_km}}

The grid and projection are shared with app.js. Do not change one without the other.
"""
import io
import json
import math
import os
import sys
import zipfile

from PIL import Image, ImageDraw
from shapely.geometry import shape, LineString, MultiLineString, Point, box
from shapely.ops import unary_union, linemerge
from shapely.prepared import prep

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAW = os.path.join(ROOT, 'data', 'raw')
OUT = os.path.join(ROOT, 'assets')
os.makedirs(OUT, exist_ok=True)

W, H = 1360, 1600
LON0, LON1, LAT0, LAT1 = 6.5, 18.7, 36.5, 47.2
K_LAT = H / (LAT1 - LAT0)
K_LON = K_LAT * math.cos(math.radians(41.9))
SS = 4  # supersampling factor for the mask
BBOX = box(LON0, LAT0, LON1, LAT1)


def proj(lon, lat):
    return (lon - LON0) * K_LON, (LAT1 - lat) * K_LAT


def km(coords):
    """Length of a lon/lat polyline in km (haversine)."""
    total = 0.0
    for (a, b) in zip(coords, coords[1:]):
        la1, lo1, la2, lo2 = map(math.radians, (a[1], a[0], b[1], b[0]))
        h = math.sin((la2 - la1) / 2) ** 2 + math.cos(la1) * math.cos(la2) * math.sin((lo2 - lo1) / 2) ** 2
        total += 2 * 6371.0 * math.asin(math.sqrt(h))
    return total


def log(*a):
    print(*a, file=sys.stderr, flush=True)


# ---------------------------------------------------------------- land
def load_land():
    polys = []
    src = None
    p = os.path.join(RAW, 'gb_ITA_ADM0.geojson')
    try:
        gj = json.load(open(p))
        polys += [shape(f['geometry']) for f in gj['features']]
        src = 'geoBoundaries (OSM-derived)'
    except Exception as e:  # truncated download or missing
        log('geoBoundaries ITA unusable (%s); falling back to Natural Earth 10m' % e)
        import shapefile
        z = zipfile.ZipFile(os.path.join(RAW, 'ne_10m_admin_0_countries.zip'))
        r = shapefile.Reader(shp=io.BytesIO(z.read('ne_10m_admin_0_countries.shp')),
                             dbf=io.BytesIO(z.read('ne_10m_admin_0_countries.dbf')),
                             shx=io.BytesIO(z.read('ne_10m_admin_0_countries.shx')))
        fields = [f[0] for f in r.fields[1:]]
        for sr in r.shapeRecords():
            rec = dict(zip(fields, sr.record))
            if rec.get('ADM0_A3') in ('ITA', 'SMR', 'VAT'):
                polys.append(shape(sr.shape.__geo_interface__))
        src = 'Natural Earth 10m'
    for iso in ('SMR', 'VAT'):
        p = os.path.join(RAW, 'gb_%s_ADM0.geojson' % iso)
        try:
            gj = json.load(open(p))
            polys += [shape(f['geometry']) for f in gj['features']]
        except Exception:
            pass
    land = unary_union([g.buffer(0) for g in polys]).intersection(BBOX)
    # Administrative polygons treat lakes as territory; the mold should not swim across Garda.
    lakes_zip = os.path.join(RAW, 'ne_10m_lakes.zip')
    if os.path.exists(lakes_zip):
        import shapefile
        z = zipfile.ZipFile(lakes_zip)
        r = shapefile.Reader(shp=io.BytesIO(z.read('ne_10m_lakes.shp')),
                             dbf=io.BytesIO(z.read('ne_10m_lakes.dbf')),
                             shx=io.BytesIO(z.read('ne_10m_lakes.shx')))
        lakes = [shape(sr.shape.__geo_interface__).buffer(0) for sr in r.shapeRecords()]
        lakes = [l for l in lakes if l.intersects(BBOX) and l.area > 0.002]  # > ~20 km2
        if lakes:
            land = land.difference(unary_union(lakes))
            log('lakes subtracted:', len(lakes))
    return land, src


def rasterize_land(land):
    img = Image.new('L', (W * SS, H * SS), 0)
    d = ImageDraw.Draw(img)
    geoms = list(land.geoms) if hasattr(land, 'geoms') else [land]
    n = 0
    for g in geoms:
        if g.is_empty or g.geom_type != 'Polygon':
            continue
        ext = [tuple(v * SS for v in proj(x, y)) for x, y in g.exterior.coords]
        if len(ext) >= 3:
            d.polygon(ext, fill=255)
            n += 1
        for ring in g.interiors:
            hole = [tuple(v * SS for v in proj(x, y)) for x, y in ring.coords]
            if len(hole) >= 3:
                d.polygon(hole, fill=0)
    log('land polygons drawn:', n)
    return img.resize((W, H), Image.BOX)


# ---------------------------------------------------------------- ferries
def load_ferries(land):
    p = os.path.join(RAW, 'osm_ferry.json')
    if not os.path.exists(p):
        return []
    els = json.load(open(p)).get('elements', [])
    land_prep = prep(land)
    near = prep(land.buffer(0.03))  # ~3 km
    kept = []
    for w in els:
        geom = w.get('geometry') or []
        coords = [(g['lon'], g['lat']) for g in geom]
        if len(coords) < 2:
            continue
        length = km(coords)
        if length < 2.5:
            continue
        a, b = Point(coords[0]), Point(coords[-1])
        if not (near.contains(a) and near.contains(b)):
            continue  # international or lake-to-nowhere
        mid = LineString(coords).interpolate(0.5, normalized=True)
        if land_prep.contains(mid):
            continue  # lake or river ferry: water inside the land polygon
        # Venice lagoon vaporetti are mapped as ferries; the lagoon is water in the mask,
        # but those are boats between quays, not sea corridors.
        minx, miny, maxx, maxy = LineString(coords).bounds
        if 12.2 < minx and maxx < 12.65 and 45.15 < miny and maxy < 45.65:
            continue
        name = (w.get('tags') or {}).get('name', '')
        kept.append({'coords': coords, 'km': length, 'name': name})
    kept.sort(key=lambda k: -k['km'])
    log('ferry corridors kept: %d of %d' % (len(kept), len(els)))
    for k in kept[:40]:
        log('   %6.1f km  %s' % (k['km'], k['name']))
    return kept


def rasterize_ferries(ferries):
    img = Image.new('L', (W * SS, H * SS), 0)
    d = ImageDraw.Draw(img)
    for f in ferries:
        pts = [tuple(v * SS for v in proj(x, y)) for x, y in f['coords']]
        d.line(pts, fill=255, width=3 * SS, joint='curve')
    return img.resize((W, H), Image.BOX)


# ---------------------------------------------------------------- towns
EXONYMS = {
    'Rome': 'Roma', 'Milan': 'Milano', 'Naples': 'Napoli', 'Turin': 'Torino', 'Florence': 'Firenze',
    'Venice': 'Venezia', 'Genoa': 'Genova', 'Padua': 'Padova', 'Syracuse': 'Siracusa', 'Mantua': 'Mantova',
    'Leghorn': 'Livorno', 'Sassari': 'Sassari', 'Bolzano': 'Bolzano', 'Trent': 'Trento', 'Lucca': 'Lucca',
}


def load_cities():
    z = zipfile.ZipFile(os.path.join(RAW, 'cities5000.zip'))
    rows = []
    with z.open('cities5000.txt') as f:
        for line in io.TextIOWrapper(f, encoding='utf-8'):
            c = line.rstrip('\n').split('\t')
            if c[8] not in ('IT', 'SM', 'VA') or c[6] != 'P':
                continue
            lat, lon, pop = float(c[4]), float(c[5]), int(c[14] or 0)
            if pop <= 0 or not (LON0 <= lon <= LON1 and LAT0 <= lat <= LAT1):
                continue
            name = EXONYMS.get(c[1], c[1])
            rows.append([name, round(lat, 4), round(lon, 4), pop])
    rows.sort(key=lambda r: -r[3])
    log('towns kept: %d (>=15k: %d, >=50k: %d)' % (len(rows), sum(r[3] >= 15000 for r in rows), sum(r[3] >= 50000 for r in rows)))
    return rows


# ---------------------------------------------------------------- overlays
def pack(geoms, tol, nd):
    """Merge touching ways into long polylines, simplify, round. Keeps the JSON small."""
    merged = linemerge([g for g in geoms if g.geom_type == 'LineString' and len(g.coords) >= 2])
    parts = list(merged.geoms) if hasattr(merged, 'geoms') else [merged]
    out = []
    for part in parts:
        ls = part.simplify(tol, preserve_topology=False)
        if len(ls.coords) >= 2:
            out.append([[round(x, nd), round(y, nd)] for x, y in ls.coords])
    return out


def osm_lines(fname, land, tol=0.0008, nd=4):
    """OSM ways from an Overpass bbox query, clipped to Italian territory (the bbox also catches
    Swiss, Austrian, French and Slovenian track)."""
    p = os.path.join(RAW, fname)
    if not os.path.exists(p):
        return None
    els = json.load(open(p)).get('elements', [])
    clip = land.buffer(0.015)
    clip_prep = prep(clip)
    geoms, total, n = [], 0.0, 0
    for w in els:
        coords = [(g['lon'], g['lat']) for g in (w.get('geometry') or [])]
        if len(coords) < 2:
            continue
        ls = LineString(coords)
        if not clip_prep.intersects(ls):
            continue
        if not clip_prep.contains(ls):
            ls = ls.intersection(clip)
        parts = list(ls.geoms) if hasattr(ls, 'geoms') else [ls]
        for part in parts:
            if part.geom_type == 'LineString' and len(part.coords) >= 2:
                total += km(list(part.coords))
                geoms.append(part)
        n += 1
    lines = pack(geoms, tol, nd)
    log('%s: %d ways -> %d polylines, %.0f km' % (fname, n, len(lines), total))
    return {'lines': lines, 'km': round(total)}


def roman_lines(land, tol=0.0015):
    p = os.path.join(RAW, 'itinere_roads.geojson')
    if not os.path.exists(p):
        return None
    try:
        gj = json.load(open(p))
    except Exception as e:
        log('itinere unusable:', e)
        return None
    feats = gj['features']
    # Itiner-e ships in EPSG:3395 (World Mercator, metres); convert to lon/lat.
    R = 6378137.0

    def to_lonlat(x, y):
        return math.degrees(x / R), math.degrees(2 * math.atan(math.exp(y / R)) - math.pi / 2)

    land_prep = prep(land.buffer(0.01))
    geoms, total, kept, main = [], 0.0, 0, 0
    for f in feats:
        g = f.get('geometry')
        if not g:
            continue
        props = f.get('properties') or {}
        if props.get('Type') != 'Main Road':   # viae publicae only; secondary roads triple the size
            continue
        main += 1
        parts = g['coordinates'] if g['type'] == 'MultiLineString' else [g['coordinates']]
        for part in parts:
            coords = [to_lonlat(x, y) for x, y, *_ in part]
            if len(coords) < 2:
                continue
            ls = LineString(coords)
            if not ls.intersects(BBOX) or not land_prep.intersects(ls):
                continue
            total += km(coords)
            geoms.append(ls)
            kept += 1
    lines = pack(geoms, tol, 3)
    log('itinere: %d main-road features, %d segments touch Italy -> %d polylines, %.0f km' % (main, kept, len(lines), total))
    return {'lines': lines, 'km': round(total)}


def main():
    land, src = load_land()
    log('land source:', src, '| polygons:', len(land.geoms) if hasattr(land, 'geoms') else 1)
    landimg = rasterize_land(land)
    ferries = load_ferries(land)
    ferryimg = rasterize_ferries(ferries)
    mask = Image.merge('RGB', (landimg, ferryimg, Image.new('L', (W, H), 0)))
    mask.save(os.path.join(OUT, 'mask.png'), optimize=True)
    land_cells = sum(1 for v in landimg.getdata() if v >= 128)
    log('mask.png written: %dx%d, land cells %d (%.1f%%)' % (W, H, land_cells, 100.0 * land_cells / (W * H)))

    cities = load_cities()
    json.dump(cities, open(os.path.join(OUT, 'cities.json'), 'w'), ensure_ascii=False, separators=(',', ':'))

    overlays = {}
    for key, fname, tol, nd in (('hsr', 'osm_hsr.json', 0.0008, 4), ('rail', 'osm_rail_main.json', 0.0015, 3), ('motorway', 'osm_motorway.json', 0.0015, 3)):
        o = osm_lines(fname, land, tol, nd)
        if o:
            overlays[key] = o
    if ferries:
        overlays['ferry'] = {'lines': [[[round(x, 4), round(y, 4)] for x, y in f['coords']] for f in ferries],
                             'km': round(sum(f['km'] for f in ferries))}
    r = roman_lines(land)
    if r:
        overlays['roman'] = r
    json.dump(overlays, open(os.path.join(OUT, 'overlays.json'), 'w'), separators=(',', ':'))
    for k, v in overlays.items():
        log('overlay %-8s %5d lines %7d km' % (k, len(v['lines']), v['km']))
    summary = {'land_source': src, 'towns': len(cities), 'ferries': len(ferries), 'overlays': {k: {'lines': len(v['lines']), 'km': v['km']} for k, v in overlays.items()}}
    json.dump(summary, open(os.path.join(ROOT, 'data', 'summary.json'), 'w'), indent=1)
    for f in ('mask.png', 'cities.json', 'overlays.json'):
        log('%-14s %8.1f KB' % (f, os.path.getsize(os.path.join(OUT, f)) / 1024))


if __name__ == '__main__':
    main()