"""Capas adicionales para el modelo de estrellas.

Pasada A (líneas y nodos): cursos de agua y pistas clasificadas por calidad.
Pasada B (áreas): arbolado, roca desnuda, matorral y cultivo.

La calidad de la pista importa tanto como su existencia: una grade1 compactada
admite una furgoneta cargada, una grade5 embarrada en un repecho no. OSM lo
etiqueta con tracktype y surface, así que se separan en dos capas en vez de
tratar todas las pistas por igual como hacía el primer modelo.
"""
import pickle
import sys
import time
from array import array

import numpy as np
import osmium
import osmium.geom
from shapely import wkb as shapely_wkb
from shapely.ops import transform as shp_transform

sys.path.insert(0, str(__file__.rsplit("/", 1)[0]))
import config as C
from extract_osm import Lines, Pts  # reutilizamos los acumuladores

WKB = osmium.geom.WKBFactory()

WATER_LINE = {"river", "stream", "canal"}
GOOD_TRACK = {"grade1", "grade2"}
GOOD_SURF = {"paved", "asphalt", "concrete", "compacted", "gravel",
             "fine_gravel", "pebblestone"}
BAD_SURF = {"mud", "sand", "grass", "ground", "dirt", "earth"}


def area_class(tags):
    n, lu = tags.get("natural"), tags.get("landuse")
    if n == "wood" or lu == "forest":
        return "forest"
    if n in ("bare_rock", "scree", "shingle", "cliff"):
        return "rock"
    if n in ("scrub", "heath", "grassland"):
        return "scrub"
    if lu in ("farmland", "orchard", "vineyard"):
        return "farmland"
    return None


def pass_lines(pbf, suffix):
    water = Lines()
    track_good, track_bad = Lines(), Lines()
    springs = Pts()

    keys = osmium.filter.KeyFilter("waterway", "highway", "natural")
    fp = (osmium.FileProcessor(pbf, osmium.osm.NODE | osmium.osm.WAY)
          .with_locations("flex_mem").with_filter(keys))

    t0 = time.time()
    for obj in fp:
        tags = obj.tags
        if obj.is_node():
            if tags.get("natural") == "spring" and obj.location.valid():
                springs.add(obj.location.lon, obj.location.lat)
            continue

        ww, hw = tags.get("waterway"), tags.get("highway")
        if ww is None and hw != "track":
            continue
        try:
            lons = [nd.lon for nd in obj.nodes]
            lats = [nd.lat for nd in obj.nodes]
        except osmium.InvalidLocationError:
            continue
        if len(lons) < 2:
            continue

        if ww in WATER_LINE:
            water.add(lons, lats)
        elif hw == "track":
            tt, sf = tags.get("tracktype"), tags.get("surface")
            good = (tt in GOOD_TRACK) or (sf in GOOD_SURF)
            bad = (sf in BAD_SURF) or (tt in ("grade4", "grade5"))
            (track_bad if (bad and not good) else track_good).add(lons, lats)

    print(f"  pasada de líneas {time.time()-t0:.0f}s", flush=True)

    def save(name, lon, lat):
        lon, lat = np.asarray(lon), np.asarray(lat)
        m = ((lon >= C.LON_MIN) & (lon <= C.LON_MAX) &
             (lat >= C.LAT_MIN) & (lat <= C.LAT_MAX))
        x, y = C.to_grid(lon[m], lat[m])
        np.save(C.INTERIM / f"{name}{suffix}.npy",
                np.stack([x, y]).astype(np.float32))
        print(f"    {name+suffix:20s} {int(m.sum()):>12,} pts", flush=True)

    save("line_water", *water.densify())
    save("line_track_good", *track_good.densify())
    save("line_track_bad", *track_bad.densify())
    save("springs", *springs.arrays())


def pass_areas(pbf, suffix):
    keys = osmium.filter.KeyFilter("natural", "landuse")
    fp = osmium.FileProcessor(pbf).with_areas().with_filter(keys)
    out = {"forest": [], "rock": [], "scrub": [], "farmland": []}
    t0 = time.time()
    n = 0
    for obj in fp:
        if not obj.is_area():
            continue
        cat = area_class(obj.tags)
        if cat is None:
            continue
        try:
            g = shapely_wkb.loads(WKB.create_multipolygon(obj), hex=True)
        except Exception:
            continue
        if g.is_empty:
            continue
        n += 1
        out[cat].append(g)
    print(f"  pasada de áreas: {n:,} en {time.time()-t0:.0f}s", flush=True)

    for cat, geoms in out.items():
        recs = []
        for g in geoms:
            try:
                recs.append(shp_transform(lambda xx, yy: C.to_grid(xx, yy), g))
            except Exception:
                continue
        with open(C.INTERIM / f"area_{cat}{suffix}.pkl", "wb") as fh:
            pickle.dump(recs, fh, protocol=4)
        print(f"    area_{cat+suffix:18s} {len(recs):>10,}", flush=True)


if __name__ == "__main__":
    pbf = sys.argv[1] if len(sys.argv) > 1 else str(C.PBF)
    suf = sys.argv[2] if len(sys.argv) > 2 else ""
    print(f"== {pbf} ==", flush=True)
    pass_lines(pbf, suf)
    pass_areas(pbf, suf)