#!/usr/bin/python3

import sys
import gzip
import time
import datetime
import math
from lxml import etree
# Knihovna lxml je Python binding na C knihovnu lxml, takže většina zpracování XMLka
# je implementovaná rychle v C.
# Lze ji získat pomocí `pip3 install lxml`


class Area:
    def __init__(self):
        self.borderWays = {}
        self.borderNodes = {}

    def computeBoundary(self):
        """Funkce volaná na oblasti pro spočítání její min/max lon/lat"""
        self.lonMin, self.lonMax = math.inf, -math.inf
        self.latMin, self.latMax = math.inf, -math.inf
        for wayID in self.borderWays:
            for node in self.borderWays[wayID]:
                self.latMin = min(self.latMin, node.lat)
                self.latMax = max(self.latMax, node.lat)
                self.lonMin = min(self.lonMin, node.lon)
                self.lonMax = max(self.lonMax, node.lon)

    def isInside(self, point):
        """Funkce volaná na oblasti pro zjištění, jestli je zadaný bod uvnitř, nebo ne"""
        if point.lat < self.latMin or point.lat > self.latMax or point.lon < self.lonMin or point.lon > self.lonMax:
            # Je mimo obdélník, nemá smysl kontrolovat podrobněji
            return False

        # Kontrolu uděláme pomocí klasického přístupu s polopřímkou vypálenou
        # náhodným směrem od bodu, budeme počítat počet průniků s hranami.
        # Pro zjednodušení nám stačí vzít horizontální polopřímku.
        horizontalLine = point.lat
        counter = 0
        for wayID in self.borderWays:
            way = self.borderWays[wayID]
            for i in range(1, len(way)):
                # Pro eliminaci chyb při průchodu horizontální čáry přímo koncem hrany
                # potřebujeme, aby všechny hrany vedly vždy stejně. Nechť je to "zdola nahoru".
                a, b = way[i-1], way[i]
                if a.lat > b.lat:
                    a, b = b, a

                # Protínat se mohou jenom pokud je jeden bod hrany pod (včetně) a druhý nad linkou (nevčetně)
                if (a.lat <= horizontalLine) and (b.lat > horizontalLine):
                    # Spočítáme bod, kde se hrana protíná s horizontální přímkou
                    # 1. Vezmeme si vektor (a-b) a spočítáme, ve které jeho části protne horizontalLine
                    t = (a.lat - horizontalLine) / (a.lat - b.lat)  # t by mělo vyjít mezi 0 a 1

                    # 2. Spočítáme longtitude v bodě protnutí
                    crossLon = a.lon + t*(b.lon-a.lon)

                    # 3. Pokud se protínají na polopřímce doprava od zkoumaného bodu, tak započítáme
                    if crossLon >= point.lon:
                        counter += 1

        # Pokud je počet protnutí lichý -> je v oblasti, jinak ne
        return (counter % 2 == 1)


class Node:
    def __init__(self, id):
        self.id = id
        self.lat = math.nan
        self.lon = math.nan

    def setLatLon(self, lat, lon):
        self.lat = lat
        self.lon = lon


################################################################################

def die(exit_code, message):
    print(message, file=sys.stderr)
    sys.exit(exit_code)


def hasTagValue(el, name, value):
    """Kontroluje jestli má element daný tag s danou hodnotou.

    Arguments:
        el {etree.Element} -- Naparsovaný XML element
        name {str} -- Název tagu
        value {str} -- Hodnota ke kontrole

    Returns:
        bool
    """
    for tag in el.findall("tag"):
        if tag.get('k') == name:
            return tag.get('v') == value
    return False


start = time.time()
phaseStart = start


# Funkce na vypisování statistiky
def printStat(parsed):
    megabytes = parsed / 1024 / 1024
    elapsed = time.time() - phaseStart
    elapsedString = str(datetime.timedelta(seconds=elapsed))
    speed = megabytes/elapsed if elapsed > 0 else 0
    print(
        f"Elapsed: {elapsedString:20s} Read: {megabytes:10.2f}MB ({speed:5.2f}MB/s)",
        file=sys.stderr, end="\r", flush=True
    )

################################################################################


if len(sys.argv) != 3:
    die(1, f"Usage: {sys.argv[0]} <file.xml.gz> <area_id>")

areaID = int(sys.argv[2])

################################################################################
# PRVNÍ PRŮCHOD - získání relace zadávající oblast #############################

print("[Phase 1] Getting relation with given ID to determine which ways are needed", file=sys.stderr)

area = Area()

with gzip.open(sys.argv[1], 'rb') as gzippedFile:
    # Budeme streamově parsovat pomocí iterparse(). Ten se může "chytat" na více
    # typů eventů - nám bude stačit defaultní "end", který nastává vždy po načtení
    # celého tagu (tedy třeba po načtení </node>). Je potřeba vyjmenovat elementy,
    # které nás zajímají, abychom nebyli voláni například pro vnořené <tag> elementy
    # a nemuseli to speciálně ošetřovat.
    i = 0
    found = False
    for event, elem in etree.iterparse(gzippedFile, tag=('node', 'way', 'relation')):
        if i % 10000 == 0:
            printStat(gzippedFile.tell())
        i += 1
        if elem.tag == 'relation':
            ID = int(elem.get('id'))
            if ID == areaID:
                for member in elem.findall("member"):
                    if member.get("type") == "way" and member.get("role") != "inner":
                        ID = int(member.get("ref"))
                        area.borderWays[ID] = []
                found = True
                break  # Není potřeba tento průchod dále protahovat

        # Odstraníme zpracovaný element z paměti (musíme udělat explicitně, jinak by tam zůstal).
        # Musíme smazat i všechny rodiče současného elementu, kteří byli vytvořeni při parsování
        # (jinak by nám elementy zůstaly a postupně užíraly paměť, souvisí to s interním fungováním
        # knihovny lxml)
        elem.clear()
        while elem.getprevious() is not None:
            del elem.getparent()[0]
        del elem

    if not found:
        die(2, "\nArea with ID '%s' not found\n", areaID)
    printStat(gzippedFile.tell())
    print(f"\nArea with {len(area.borderWays)} ways found", file=sys.stderr)

################################################################################
# DRUHÝ PRŮCHOD ################################################################

print("[Phase 2] Getting node IDs of all ways of the area", file=sys.stderr)

# Poznamenáme si ID všech nodů, které chceme (mimo uložení do cest se
# nám hodí i seznam všech nodů, ve které se dá rychle hledat)
phaseStart = time.time()
with gzip.open(sys.argv[1], 'rb') as gzippedFile:
    i = 0
    # Budeme streamově parsovat pomocí iterparse(), viz první průchod.
    for event, elem in etree.iterparse(gzippedFile, tag=('node', 'way', 'relation')):
        if i % 10000 == 0:
            printStat(gzippedFile.tell())
        i += 1
        if elem.tag == 'way':
            ID = int(elem.get("id"))
            if ID in area.borderWays:
                for nd in elem.findall("nd"):
                    nodeID = int(nd.get("ref"))
                    # Každý node je reprezentován jedním objektem, z cesty i ze
                    # seznamu nodů je to jenom pointer na tento (společný) objekt
                    if nodeID not in area.borderNodes:
                        area.borderNodes[nodeID] = Node(nodeID)
                    area.borderWays[ID].append(area.borderNodes[nodeID])

        # Odstraníme zpracovaný element z paměti (viz první průchod)
        elem.clear()
        while elem.getprevious() is not None:
            del elem.getparent()[0]
        del elem

    printStat(gzippedFile.tell())
    print(f"\nFound {len(area.borderNodes)} nodes around the area", file=sys.stderr)

################################################################################
# TŘETÍ PRŮCHOD ################################################################

print("[Phase 3] Getting all nodes around the area", file=sys.stderr)

# 1. Projdeme všechny nodes a zaznamenáme si ty, které chceme
foundCounter = 0
phaseStart = time.time()
with gzip.open(sys.argv[1], 'rb') as gzippedFile:
    i = 0
    # Budeme streamově parsovat pomocí iterparse(), viz první průchod.
    for event, elem in etree.iterparse(gzippedFile, tag=('node')):
        if i % 10000 == 0:
            printStat(gzippedFile.tell())
        i += 1
        ID = int(elem.get("id"))
        if ID in area.borderNodes:
            lat = float(elem.get("lat"))
            lon = float(elem.get("lon"))
            area.borderNodes[ID].setLatLon(lat, lon)
            foundCounter += 1
            if foundCounter == len(area.borderNodes):
                break

        # Odstraníme zpracovaný element z paměti (viz první průchod)
        elem.clear()
        while elem.getprevious() is not None:
            del elem.getparent()[0]
        del elem

    printStat(gzippedFile.tell())

# 2. Spočítáme min/max lat/long
area.computeBoundary()

print(f"\nCoordinates of {foundCounter} nodes of area found", file=sys.stderr)

################################################################################
# ČTVRTÝ PRŮCHOD ###############################################################

print("[Phase 4] Getting all wanted nodes and comparing them with parsed area", file=sys.stderr)

# 1. Projdeme všechny nodes a zaznamenáme si ty, které jsou v oblasti a mají patřičný tag
foundCounter = 0
phaseStart = time.time()
with gzip.open(sys.argv[1], 'rb') as gzippedFile:
    i = 0
    # Budeme streamově parsovat pomocí iterparse(), viz první průchod.
    for event, elem in etree.iterparse(gzippedFile, tag=('node')):
        if i % 10000 == 0:
            printStat(gzippedFile.tell())
        i += 1
        if hasTagValue(elem, "amenity", "drinking_water") or hasTagValue(elem, "drinking_water", "yes"):
            lat = float(elem.get("lat"))
            lon = float(elem.get("lon"))
            # Podrobnější kontrola
            point = Node(-1)
            point.setLatLon(lat, lon)
            if area.isInside(point):
                print(elem.get("id"))
                foundCounter += 1

        # Odstraníme zpracovaný element z paměti (viz první průchod)
        elem.clear()
        while elem.getprevious() is not None:
            del elem.getparent()[0]
        del elem

    printStat(gzippedFile.tell())

print(f"\nFound {foundCounter} 'amenity=drinkig_water' or 'drinking_water=yes' nodes", file=sys.stderr)