#!/usr/bin/env python3
"""Rebuild data/geo/SE/*.tsv from SCB (CC0), GeoNames (CC BY 4.0) and Trafikverket NVDB (CC0).

    TRAFIKVERKET_API_KEY=… data-import/geo-se.py [--key-file FILE] [--cache DIR] [--streets-per-locality N] [--out DIR]
"""
import collections
import csv
import io
import json
import os
import re
import sys
import urllib.request
import zipfile
from pathlib import Path

import geo
import source
import tsv
import xlsx

CODES = "https://www.scb.se/contentassets/7a89e48960f741e08918e489ea36354a/kommunlankod-2026.xlsx"
POPULATION = "https://api.scb.se/OV0104/v1/doris/sv/ssd/START/BE/BE0101/BE0101A/BefolkningNy"
POPULATION_QUERY = {
    "query": [
        {"code": "Region", "selection": {"filter": "all", "values": ["*"]}},
        {"code": "ContentsCode", "selection": {"filter": "item", "values": ["BE0101N1"]}},
        {"code": "Tid", "selection": {"filter": "top", "values": ["1"]}},
    ],
    "response": {"format": "json"},
}
TATORTER = "https://geodata.scb.se/geoserver/stat/wfs?service=WFS&version=2.0.0&request=GetFeature&typeNames=stat:Tatorter_2023&outputFormat=csv&propertyName=tatort,kommun,bef"
POSTAL_CODES = "https://download.geonames.org/export/zip/SE.zip"
NVDB = "https://api.trafikinfo.trafikverket.se/v2/data.json"
NVDB_PAGE = 50000
TIMEZONE = "Europe/Stockholm"
# This script's reading of the "exceptions in big cities" in docs/research/research-geo-se.md, Postal codes → Structure.
ONE_POSITION = {"Stockholm", "Göteborg", "Malmö"}
# A tätort's floor: "minst 200 invånare", docs/research/research-geo-se.md, Lantmäteriet → Ortnamn.
UNMATCHED_POPULATION = 200


def scb_codes(cache):
    regions, municipalities = {}, {}
    for cells in xlsx.rows(source.fetch(CODES, cache, "kommunlankod.xlsx", magic=b"PK")):
        if len(cells) < 2 or not re.fullmatch(r"\d{2}|\d{4}", cells[0]):
            continue
        (regions if len(cells[0]) == 2 else municipalities)[cells[0]] = cells[1].strip()
    return regions, municipalities


def scb_population(cache):
    body = json.dumps(POPULATION_QUERY).encode()
    data = source.fetch(POPULATION, cache, "befolkning.json", data=body, headers={"Content-Type": "application/json"})
    return {row["key"][0]: row["values"][0] for row in json.loads(data.decode("utf-8-sig"))["data"]}


def scb_tatorter(cache):
    text = source.fetch(TATORTER, cache, "tatorter.csv").decode("utf-8")
    by_name = collections.defaultdict(list)
    for r in csv.DictReader(io.StringIO(text)):
        by_name[r["tatort"]].append((r["kommun"], int(r["bef"])))
    return by_name


def geonames(cache):
    z = zipfile.ZipFile(io.BytesIO(source.fetch(POSTAL_CODES, cache, "SE.zip", magic=b"PK")))
    rows = []
    for line in z.read("SE.txt").decode("utf-8").splitlines():
        f = line.split("\t")
        lat, lon = (float(f[9]), float(f[10])) if f[9] and f[10] else (None, None)
        rows.append({"code": f[1], "locality": f[2], "municipality": f[6], "lat": lat, "lon": lon})
    return rows


def nvdb_segments(cache, key):
    path = cache / "nvdb-gatunamn.tsv"
    if not path.exists():
        with open(path.with_suffix(".part"), "w", encoding="utf-8") as out:
            change = "0"
            while True:
                query = (
                    f'<REQUEST><LOGIN authenticationkey="{key}"/>'
                    f'<QUERY objecttype="Gatunamn" namespace="vägdata.nvdb_dk_o" schemaversion="1.0" limit="{NVDB_PAGE}" changeid="{change}">'
                    "<FILTER><EQ name=\"Deleted\" value=\"false\"/></FILTER>"
                    "<INCLUDE>Namn</INCLUDE><INCLUDE>Geometry.WKT-WGS84-3D</INCLUDE></QUERY></REQUEST>"
                )
                req = urllib.request.Request(NVDB, data=query.encode(), headers={"Content-Type": "text/xml"})
                with urllib.request.urlopen(req, timeout=600) as r:
                    result = json.load(r)["RESPONSE"]["RESULT"][0]
                rows = result.get("Gatunamn", [])
                for row in rows:
                    m = re.match(r"LINESTRING Z \(([-\d.]+) ([-\d.]+) ", row.get("Geometry", {}).get("WKT-WGS84-3D", ""))
                    name = " ".join(row.get("Namn", "").split())
                    if m and name:
                        out.write(f"{name}\t{m.group(1)}\t{m.group(2)}\n")
                change = result["INFO"]["LASTCHANGEID"]
                if len(rows) < NVDB_PAGE:
                    break
        path.with_suffix(".part").rename(path)
    for line in path.read_text(encoding="utf-8").splitlines():
        name, lon, lat = line.split("\t")
        yield name, float(lat), float(lon)


def street_delivery(name, codes):
    """The codes delivered to a street. A postort's own prefix is three digits where one three-digit prefix starts at
    least half its codes, else two. Outside ONE_POSITION's cities, the digit after that prefix is 0 or 1 for boxes and 8
    for reply mail; inside them, a second digit 0 marks boxes. The box and reply-mail digits are per
    docs/research/research-geo-se.md, Postal codes → Structure."""
    if name in ONE_POSITION:
        return [c for c in codes if c[1] != "0"]
    largest = collections.Counter(c[:3] for c in codes).most_common(1)[0][1]
    length = 3 if largest * 2 >= len(codes) else 2
    return [c for c in codes if c[length] not in "018"]


def municipality_of(name, rows, tatorter, municipalities):
    named = [code for code, n in municipalities.items() if n == name]
    if named:
        return named[0], "kommun"
    voted = collections.Counter(r["municipality"] for r in rows if r["municipality"] in municipalities)
    matches = tatorter.get(name, [])
    if matches:
        in_vote = [m for m in matches if voted and m[0] == voted.most_common(1)[0][0]]
        return max(in_vote or matches, key=lambda m: m[1])[0], "tatort"
    if voted:
        return voted.most_common(1)[0][0], "codes"
    return None, "unplaced"


def population_of(name, municipality, tatorter, municipalities, population):
    matches = [m for m in tatorter.get(name, []) if m[0] == municipality]
    if matches:
        return str(max(m[1] for m in matches))
    return population[municipality] if municipalities[municipality] == name else str(UNMATCHED_POPULATION)


def well_cased(name):
    """Whether each word of name is capitalised: GeoNames also spells some postorter, such as STOCKHOLM and örebro, in
    capitals or lower case."""
    return all(part[:1].isupper() and (len(part) == 1 or not part.isupper()) for part in re.split(r"[ -]", name))


def localities(codes, tatorter, municipalities, population):
    """Each postort with its municipality, weight, centroid and street-delivery codes."""
    by_locality = collections.defaultdict(list)
    for r in codes:
        by_locality[r["locality"]].append(r)
    out, how = {}, collections.Counter()
    for name, rows in by_locality.items():
        municipality, method = municipality_of(name, rows, tatorter, municipalities)
        how[method] += 1
        kept = street_delivery(name, [r["code"].replace(" ", "") for r in rows])
        point = geo.centroid(rows)
        if municipality is None or not kept or point is None or not well_cased(name):
            continue
        lat, lon = point
        out[name] = {"name": name, "municipality": municipality, "population": population_of(name, municipality, tatorter, municipalities, population), "lat": f"{lat:.4f}", "lon": f"{lon:.4f}", "codes": kept}
    print(f"municipality by {dict(how)}; {len(by_locality) - len(out)} postorter dropped before streets", file=sys.stderr)
    return out


def streets(segments, codes, localities, per_locality):
    """The names with most segments per locality, each segment at its nearest code centroid."""
    nearest = geo.Nearest((r["lat"], r["lon"], r["locality"]) for r in codes if r["lat"] is not None and r["locality"] in localities)
    count = collections.Counter()
    for name, lat, lon in segments:
        if name[0].isalpha():
            count[(nearest.find(lat, lon), name)] += 1
    return geo.top_streets(count, per_locality, "segments")


def main():
    a = geo.parser(__doc__, "SE", key_file={"help": "file holding the Trafikverket API key; TRAFIKVERKET_API_KEY otherwise"}).parse_args()
    key = Path(a.key_file).read_text().strip() if a.key_file else os.environ.get("TRAFIKVERKET_API_KEY")
    if not key:
        sys.exit("set TRAFIKVERKET_API_KEY or pass --key-file")
    cache, out = geo.directories(a)

    region_names, municipality_names = scb_codes(cache)
    population = scb_population(cache)
    codes = geonames(cache)
    places = localities(codes, scb_tatorter(cache), municipality_names, population)
    named = streets(nvdb_segments(cache, key), codes, places, a.streets_per_locality)
    regions = {c: {"code": c, "name": n, "population": population[c], "timezone": TIMEZONE} for c, n in region_names.items()}
    municipalities = {c: {"code": c, "name": n, "region": c[:2], "population": population[c]} for c, n in municipality_names.items()}
    kept, dropped = geo.fence(regions, municipalities, places, named, {k: l["codes"] for k, l in places.items()})
    geo.log_dropped_localities(dropped["locality"])
    if dropped["municipality"] or dropped["region"]:
        sys.exit(f"municipalities without a locality: {sorted(dropped['municipality'])}; regions without a municipality: {sorted(dropped['region'])}")
    places = kept["locality"]

    tsv.write(out / "region.tsv", ["code", "name", "population", "timezone"], [r for _, r in sorted(kept["region"].items())])
    tsv.write(out / "municipality.tsv", ["code", "name", "region", "population"], [m for _, m in sorted(kept["municipality"].items())])
    tsv.write(out / "locality.tsv", ["name", "municipality", "population", "lat", "lon"], [l for _, l in sorted(places.items())])
    tsv.write(out / "postal-code.tsv", ["code", "locality"], sorted(({"code": f"{c[:3]} {c[3:]}", "locality": l["name"]} for l in places.values() for c in l["codes"]), key=lambda r: r["code"]))
    tsv.write(out / "street.tsv", ["name", "locality", "segments"], [s for locality in sorted(places) for s in named[locality]])


if __name__ == "__main__":
    main()