"""
Solutions for OSM Data Query and Visualization Exercises
=====================================================

This script contains the solutions for the exercises in the worksheet.
It is divided into several sections:
1. OSM Data Query and Visualization
2. Address Geocoding
"""

# =============================================================================
# 1. OSM Data Query and Visualization
# =============================================================================

# Import required libraries
import plotly.express as px  # For interactive maps
import geopandas as gpd     # For working with geodata
import matplotlib.pyplot as plt  # For static maps
from collections import namedtuple  # For structured data
from osm2geojson import overpass_call, json2geojson  # For OSM data queries

# Definition of helper structures for data organization
Bbox = namedtuple("Bbox", ["south", "west", "north", "east"])  # Bounding box for geographic areas
Tag = namedtuple("Tag", ["key", "values"])  # OSM tags for search

def load_osm_from_overpass(bbox, tag, crs="epsg:4326") -> gpd.GeoDataFrame:
    """
    Loads OSM data from the Overpass API and converts it to a GeoDataFrame.
    
    Args:
        bbox: Bounding box for the search area
        tag: OSM tag for search (e.g., sport=table_tennis)
        crs: Coordinate reference system (default: WGS84)
    
    Returns:
        GeoDataFrame with the found OSM elements
    """
    geojson = load_osm_from_overpass_geojson(bbox, tag, crs)
    return gpd.GeoDataFrame.from_features(geojson, crs=crs)

def load_osm_from_overpass_geojson(bbox, tag, crs="epsg:4326"):
    """
    Creates and executes an Overpass API query.
    
    Args:
        bbox: Bounding box for the search area
        tag: OSM tag for search
        crs: Coordinate reference system
    
    Returns:
        GeoJSON object with the found OSM elements
    """
    values = None if "*" in tag.values else tag.values
    query = f"""
    [out:json];
    node["{tag.key}"="{values[0] if values else '*'}"]({bbox.south},{bbox.west},{bbox.north},{bbox.east});
    out body;
    >;
    out skel qt;
    """
    response = overpass_call(query)
    return json2geojson(response)

def plot_gdf_on_map(gdf, color="k", title=""):
    """
    Creates an interactive map with the GeoDataFrame data.
    
    Args:
        gdf: GeoDataFrame with the data to visualize
        color: Color for display
        title: Map title
    
    Returns:
        Plotly Figure object with the interactive map
    """
    fig = px.scatter_mapbox(
        gdf,
        lat=gdf.geometry.y,
        lon=gdf.geometry.x,
        color_discrete_sequence=["fuchsia"],
        height=500,
        zoom=11,
    )
    fig.update_layout(title=title, mapbox_style="open-street-map")
    fig.update_layout(margin={"r": 0, "l": 0, "b": 0})
    return fig

def main():
    """
    Main function to test OSM data queries.
    Performs two tests:
    1. Table tennis tables in Zurich
    2. Drinking water fountains in Rapperswil-Jona
    """
    # Test 1: Table tennis tables in Zurich
    print("Test 1: Table tennis tables in Zurich")
    tag = Tag(key="sport", values=["table_tennis"])  # OSM tag for table tennis tables
    bbox_zurich = Bbox(west=8.471668, south=47.348834, east=8.600454, north=47.434379)
    table_tennis_gdf = load_osm_from_overpass(bbox_zurich, tag)
    print(f"Found table tennis tables: {len(table_tennis_gdf)}")
    
    # Test 2: Drinking water fountains in Rapperswil-Jona
    print("\nTest 2: Drinking water fountains in Rapperswil-Jona")
    tag = Tag(key="amenity", values=["drinking_water"])  # OSM tag for drinking water fountains
    bbox_rapperswil = Bbox(west=8.8000, south=47.2000, east=8.9000, north=47.3000)
    drinking_water_gdf = load_osm_from_overpass(bbox_rapperswil, tag)
    print(f"Found drinking water fountains: {len(drinking_water_gdf)}")

# =============================================================================
# 2. Address Geocoding
# =============================================================================

def geocode_address(address):
    """
    Geocodes an address using the Nominatim API.
    Converts a text address into geographic coordinates.
    
    Args:
        address: Text address (e.g., "Vulkanstrasse 106, Zürich")
    
    Returns:
        WKT string with coordinates in format 'SRID=4326;POINT(lon lat)'
        or None if no coordinates were found
    
    Note: Please respect the Nominatim API usage policy:
    https://operations.osmfoundation.org/policies/nominatim/
    """
    import requests
    from time import sleep
    
    base_url = 'https://nominatim.openstreetmap.org/search?'
    params = {
        'q': address,
        'format': 'json',
        'limit': 1,
        'User-Agent': 'OpenSchoolMaps Tutorial'  # Important: Use a specific User-Agent
    }
    
    try:
        response = requests.get(base_url, params=params)
        response.raise_for_status()  # Checks for HTTP errors
        loc = response.json()
        
        if len(loc) == 0:
            return None
        
        # Return coordinates in WKT format
        wkt = f'SRID=4326;POINT({loc[0]["lon"]} {loc[0]["lat"]})'
        sleep(1)  # Throttling according to Nominatim guidelines
        return wkt
        
    except requests.exceptions.RequestException as e:
        print(f"Error during geocoding: {e}")
        return None

def test_geocoding():
    """
    Tests the geocoding function with the example address.
    Displays the found coordinates or an error message.
    """
    address = 'Vulkanstrasse 106, Zürich'
    result = geocode_address(address)
    
    if result:
        print(f"Geocoded coordinates for '{address}':")
        print(result)
    else:
        print(f"No coordinates found for '{address}'")

# =============================================================================
# Main Program
# =============================================================================

if __name__ == "__main__":
    print("=== Starting OSM Data Queries ===")
    main()  # Performs the OSM data queries
    
    print("\n=== Testing the Geocoding Function ===")
    test_geocoding()  # Tests the geocoding function