app.py

"""
OSM 3D Environment Generator - Interactive Gradio App
Generate and visualize 3D building blocks and roads from OpenStreetMap data
"""

import gradio as gr
import requests
import math
import time
import os
from datetime import datetime
from typing import Tuple, Optional, Dict, List, Union
from dotenv import load_dotenv
import mercantile
import numpy as np
from PIL import Image
from io import BytesIO
import tempfile
import geopandas as gpd
from shapely.geometry import shape, mapping, LineString, Polygon, MultiPolygon  # if working with raw GeoJSON
from pyproj import CRS, Transformer  # optional, but part of geopandas
import json
import urllib.parse

# Constants
OVERPASS_URL = "http://overpass-api.de/api/interpreter"
MAX_RETRIES = 3
MAPBOX_TOKEN = os.getenv("MAPBOX_TOKEN", "")

def km_to_degrees(km: float, lat: float) -> Tuple[float, float]:
    """Convert kilometers to degrees at given latitude."""
    lat_offset = km / 111.0
    lon_offset = km / (111.0 * math.cos(math.radians(lat)))
    return lat_offset, lon_offset

def latlon_to_meters(lat: float, lon: float, center_lat: float, center_lon: float) -> Tuple[float, float]:
    """Convert lat/lon to meters from center point."""
    x = (lon - center_lon) * 111000 * math.cos(math.radians(center_lat))
    y = (lat - center_lat) * 111000
    return x, y

def fetch_osm_data(center_lat: float, center_lon: float, square_size_km: float) -> Optional[Dict]:
    """Fetch buildings and roads from OSM."""
    half_size_km = square_size_km / 2
    lat_offset, lon_offset = km_to_degrees(half_size_km, center_lat)
    
    bbox = {
        'south': center_lat - lat_offset,
        'north': center_lat + lat_offset,
        'west': center_lon - lon_offset,
        'east': center_lon + lon_offset
    }
    
    query = f"""
    [out:json][timeout:180];
    (
      way["building"]({bbox['south']},{bbox['west']},{bbox['north']},{bbox['east']});
      relation["building"]({bbox['south']},{bbox['west']},{bbox['north']},{bbox['east']});
      way["highway"]({bbox['south']},{bbox['west']},{bbox['north']},{bbox['east']});
    );
    out body;
    >;
    out skel qt;
    """
    
    # Retry logic
    for attempt in range(MAX_RETRIES):
        try:
            response = requests.post(OVERPASS_URL, data={'data': query}, timeout=200)
            if response.status_code == 200:
                return response.json()
            elif response.status_code in [429, 504] and attempt < MAX_RETRIES - 1:
                time.sleep((attempt + 1) * 10)
            else:
                return None
        except requests.exceptions.Timeout:
            if attempt < MAX_RETRIES - 1:
                time.sleep(10)
            else:
                return None
    return None

def get_building_height(tags: Dict) -> float:
    """Extract or estimate building height from OSM tags."""
    # Direct height tag
    if 'height' in tags:
        try:
            return float(tags['height'].replace('m', '').replace('M', '').strip())
        except:
            pass
    
    if 'building:height' in tags:
        try:
            return float(tags['building:height'].replace('m', '').replace('M', '').strip())
        except:
            pass
    
    # Calculate from levels
    if 'building:levels' in tags:
        try:
            return float(tags['building:levels']) * 3.5
        except:
            pass
    
    if 'levels' in tags:
        try:
            return float(tags['levels']) * 3.5
        except:
            pass
    
    # Height estimates by building type
    building_type = tags.get('building', 'yes').lower()
    height_map = {
        'cathedral': 30, 'church': 20, 'commercial': 12, 'office': 15,
        'hotel': 20, 'apartments': 18, 'residential': 9, 'house': 7,
        'garage': 3, 'shed': 3, 'tower': 40, 'stadium': 25
    }
    return height_map.get(building_type, 9.0)

def get_road_width(road_type: str) -> float:
    """Get road width in meters based on highway type."""
    width_map = {
        'motorway': 12.0, 'trunk': 10.0, 'primary': 8.0, 'secondary': 7.0,
        'tertiary': 6.0, 'residential': 5.0, 'service': 3.0, 'footway': 2.0,
        'path': 1.5, 'cycleway': 2.0, 'unclassified': 5.0
    }
    return width_map.get(road_type, 4.0)


# ============================================================
# SMOOTH ROAD FOOTPRINT ALGORITHM (V2)
# ============================================================

def offset_polyline(points: List[Tuple[float, float]], distance: float) -> List[Tuple[float, float]]:
    """
    Offset a polyline by perpendicular distance.
    Positive distance = left side, negative = right side.
    Uses averaged normals at interior points for smooth corners.
    """
    if len(points) < 2:
        return points
    
    result = []
    
    for i in range(len(points)):
        if i == 0:
            # First point: use direction to next
            dx = points[1][0] - points[0][0]
            dy = points[1][1] - points[0][1]
        elif i == len(points) - 1:
            # Last point: use direction from previous
            dx = points[i][0] - points[i-1][0]
            dy = points[i][1] - points[i-1][1]
        else:
            # Middle: average incoming and outgoing directions
            dx1 = points[i][0] - points[i-1][0]
            dy1 = points[i][1] - points[i-1][1]
            len1 = math.sqrt(dx1**2 + dy1**2)
            
            dx2 = points[i+1][0] - points[i][0]
            dy2 = points[i+1][1] - points[i][1]
            len2 = math.sqrt(dx2**2 + dy2**2)
            
            if len1 > 0.0001:
                dx1, dy1 = dx1/len1, dy1/len1
            if len2 > 0.0001:
                dx2, dy2 = dx2/len2, dy2/len2
            
            dx = dx1 + dx2
            dy = dy1 + dy2
        
        length = math.sqrt(dx**2 + dy**2)
        if length < 0.0001:
            result.append(points[i])
            continue
        
        dx /= length
        dy /= length
        
        # Perpendicular (90° left)
        perp_x = -dy * distance
        perp_y = dx * distance
        
        result.append((points[i][0] + perp_x, points[i][1] + perp_y))
    
    return result


def create_arc_points(center: Tuple[float, float], radius: float,
                      start_angle: float, end_angle: float,
                      segments: int = 8) -> List[Tuple[float, float]]:
    """Create points along an arc from start_angle to end_angle."""
    points = []
    
    # Normalize angle difference to take shorter path
    angle_diff = end_angle - start_angle
    while angle_diff > math.pi:
        angle_diff -= 2 * math.pi
    while angle_diff < -math.pi:
        angle_diff += 2 * math.pi
    
    for i in range(segments + 1):
        t = i / segments
        angle = start_angle + angle_diff * t
        x = center[0] + radius * math.cos(angle)
        y = center[1] + radius * math.sin(angle)
        points.append((x, y))
    
    return points


def get_direction_at_endpoint(points: List[Tuple[float, float]], at_start: bool) -> Tuple[float, float]:
    """Get normalized direction vector at start or end of polyline."""
    if at_start:
        dx = points[1][0] - points[0][0]
        dy = points[1][1] - points[0][1]
    else:
        n = len(points)
        dx = points[n-1][0] - points[n-2][0]
        dy = points[n-1][1] - points[n-2][1]
    
    length = math.sqrt(dx**2 + dy**2)
    if length > 0.0001:
        return (dx/length, dy/length)
    return (1, 0)


def find_road_intersections(roads: List[Dict], tolerance: float = 1.0) -> Dict[Tuple[float, float], List[Dict]]:
    """
    Find where road endpoints meet.
    Returns dict: intersection_center -> list of road info at that intersection.
    """
    # Collect all endpoints
    endpoints = []
    for road in roads:
        points = road['points']
        if len(points) < 2:
            continue
        endpoints.append({
            'point': points[0],
            'road': road,
            'position': 'start',
            'direction': get_direction_at_endpoint(points, True)
        })
        endpoints.append({
            'point': points[-1],
            'road': road,
            'position': 'end',
            'direction': get_direction_at_endpoint(points, False)
        })
    
    # Group nearby endpoints
    intersections = {}
    used = set()
    
    for i, ep1 in enumerate(endpoints):
        if i in used:
            continue
        
        group = [ep1]
        center_x = ep1['point'][0]
        center_y = ep1['point'][1]
        
        for j, ep2 in enumerate(endpoints):
            if j <= i or j in used:
                continue
            
            dist = math.sqrt((ep1['point'][0] - ep2['point'][0])**2 + 
                           (ep1['point'][1] - ep2['point'][1])**2)
            
            if dist < tolerance:
                group.append(ep2)
                center_x += ep2['point'][0]
                center_y += ep2['point'][1]
                used.add(j)
        
        if len(group) >= 2:
            center = (center_x / len(group), center_y / len(group))
            intersections[center] = group
            used.add(i)
    
    return intersections


class RoadMeshBuilder:
    """Builds proper mesh geometry for smooth road footprints."""
    
    def __init__(self, base_z: float = 0.1):
        self.vertices = []
        self.faces = []
        self.z_height = base_z
    
    def add_vertex(self, x: float, y: float, z: float = None) -> int:
        """Add a vertex and return its index (1-based for OBJ)."""
        idx = len(self.vertices) + 1  # OBJ is 1-indexed
        if z is None:
            z = self.z_height
        self.vertices.append((x, y, z))
        return idx
    
    def add_triangle(self, i0: int, i1: int, i2: int):
        """Add a triangle face (indices are 1-based)."""
        self.faces.append((i0, i1, i2))
    
    def add_quad_strip(self, left_points: List[Tuple[float, float]], 
                       right_points: List[Tuple[float, float]],
                       z_values_left: List[float] = None,
                       z_values_right: List[float] = None):
        """
        Add geometry for a quad strip between two parallel edge lines.
        This creates proper road ribbon triangulation.
        """
        n = min(len(left_points), len(right_points))
        if n < 2:
            return
        
        # Add all vertices first
        left_indices = []
        right_indices = []
        
        for i in range(n):
            z_left = z_values_left[i] if z_values_left else self.z_height
            z_right = z_values_right[i] if z_values_right else self.z_height
            
            left_indices.append(self.add_vertex(left_points[i][0], left_points[i][1], z_left))
            right_indices.append(self.add_vertex(right_points[i][0], right_points[i][1], z_right))
        
        # Create triangles for each quad
        for i in range(n - 1):
            li = left_indices[i]
            li1 = left_indices[i + 1]
            ri = right_indices[i]
            ri1 = right_indices[i + 1]
            
            # Two triangles per quad
            self.add_triangle(li, li1, ri1)
            self.add_triangle(li, ri1, ri)
    
    def add_semicircle_cap(self, center: Tuple[float, float], radius: float,
                           direction: Tuple[float, float], z: float = None):
        """Add a semicircle end cap at dead ends."""
        if z is None:
            z = self.z_height
            
        dir_angle = math.atan2(direction[1], direction[0])
        start_angle = dir_angle - math.pi/2
        end_angle = dir_angle + math.pi/2
        
        arc_points = create_arc_points(center, radius, start_angle, end_angle, 8)
        
        # Fan triangulation from center
        center_idx = self.add_vertex(center[0], center[1], z)
        
        arc_indices = []
        for pt in arc_points:
            arc_indices.append(self.add_vertex(pt[0], pt[1], z))
        
        for i in range(len(arc_indices) - 1):
            self.add_triangle(center_idx, arc_indices[i], arc_indices[i + 1])
    
    def add_fillet_arc(self, center: Tuple[float, float],
                       point1: Tuple[float, float], point2: Tuple[float, float],
                       radius: float, z: float = None):
        """Add a fillet arc connecting two edge endpoints at an intersection."""
        if z is None:
            z = self.z_height
            
        angle1 = math.atan2(point1[1] - center[1], point1[0] - center[0])
        angle2 = math.atan2(point2[1] - center[1], point2[0] - center[0])
        
        arc_points = create_arc_points(center, radius, angle1, angle2, 6)
        
        if len(arc_points) < 2:
            return
        
        center_idx = self.add_vertex(center[0], center[1], z)
        
        arc_indices = []
        for pt in arc_points:
            arc_indices.append(self.add_vertex(pt[0], pt[1], z))
        
        for i in range(len(arc_indices) - 1):
            self.add_triangle(center_idx, arc_indices[i], arc_indices[i + 1])


def generate_smooth_roads(roads_data: List[Dict], elevation_data, 
                          get_elevation_func, use_terrain: bool = True) -> Tuple[List, List]:
    """
    Generate smooth road footprint geometry using the V2 algorithm.
    
    Args:
        roads_data: List of dicts with 'points' (list of (x,y) tuples), 'width', 'type', 'lat_lon_points'
        elevation_data: Elevation grid data tuple
        get_elevation_func: Function to get elevation at a lat/lon point
        use_terrain: If True, use elevation data; if False, all roads at z=0.2
    
    Returns:
        (vertices, faces) lists for OBJ export
    """
    if not roads_data:
        return [], []
    
    # Find intersections
    intersections = find_road_intersections(roads_data, tolerance=1.0)
    
    # Track which road endpoints are at intersections, and store their Z values
    intersection_endpoints = set()
    intersection_z_values = {}  # (road_id, position) -> z_value
    
    for center, group in intersections.items():
        for ep in group:
            key = (id(ep['road']), ep['position'])
            intersection_endpoints.add(key)
    
    # Base Z when terrain is disabled
    flat_z = 0.2
    
    builder = RoadMeshBuilder(base_z=flat_z)
    
    # First pass: calculate Z values for all road endpoints
    road_z_values = {}  # road_id -> {'start': z, 'end': z, 'all': [z1, z2, ...]}
    
    for road in roads_data:
        points = road['points']
        lat_lon_points = road.get('lat_lon_points', [])
        
        if len(points) < 2:
            continue
        
        # Get Z values for each point
        z_values = []
        for i, pt in enumerate(points):
            if use_terrain and elevation_data and i < len(lat_lon_points):
                lat, lon = lat_lon_points[i]
                z = get_elevation_func(lat, lon, elevation_data) + 0.2
            else:
                z = flat_z
            z_values.append(z)
        
        road_z_values[id(road)] = {
            'start': z_values[0],
            'end': z_values[-1],
            'all': z_values
        }
    
    # Process each road
    for road in roads_data:
        points = road['points']
        width = road['width']
        
        if len(points) < 2:
            continue
        
        half_width = width / 2
        
        # Create offset lines
        left_edge = offset_polyline(points, half_width)
        right_edge = offset_polyline(points, -half_width)
        
        # Get Z values
        z_info = road_z_values.get(id(road), {'all': [flat_z] * len(points)})
        z_values = z_info['all']
        
        # Add the main quad strip
        builder.add_quad_strip(left_edge, right_edge, z_values, z_values)
        
        # Check if start is a dead end
        start_key = (id(road), 'start')
        if start_key not in intersection_endpoints:
            center = points[0]
            direction = get_direction_at_endpoint(points, True)
            z = z_values[0] if z_values else flat_z
            builder.add_semicircle_cap(center, half_width, (-direction[0], -direction[1]), z)
        
        # Check if end is a dead end
        end_key = (id(road), 'end')
        if end_key not in intersection_endpoints:
            center = points[-1]
            direction = get_direction_at_endpoint(points, False)
            z = z_values[-1] if z_values else flat_z
            builder.add_semicircle_cap(center, half_width, direction, z)
    
    # Process intersections - add fillets
    for center, group in intersections.items():
        if len(group) < 2:
            continue
        
        # Calculate intersection Z as average of all connected road endpoints
        intersection_z_list = []
        for ep in group:
            road = ep['road']
            z_info = road_z_values.get(id(road), {'start': flat_z, 'end': flat_z})
            if ep['position'] == 'start':
                intersection_z_list.append(z_info['start'])
            else:
                intersection_z_list.append(z_info['end'])
        
        intersection_z = sum(intersection_z_list) / len(intersection_z_list) if intersection_z_list else flat_z
        
        # Get edge info for each road at this intersection
        road_edges = []
        for ep in group:
            road = ep['road']
            points = road['points']
            width = road['width']
            half_width = width / 2
            
            left_edge = offset_polyline(points, half_width)
            right_edge = offset_polyline(points, -half_width)
            
            if ep['position'] == 'start':
                left_pt = left_edge[0]
                right_pt = right_edge[0]
                direction = ep['direction']
            else:
                left_pt = left_edge[-1]
                right_pt = right_edge[-1]
                direction = (-ep['direction'][0], -ep['direction'][1])
            
            angle = math.atan2(-direction[1], -direction[0])
            
            road_edges.append({
                'road': road,
                'position': ep['position'],
                'left_pt': left_pt,
                'right_pt': right_pt,
                'angle': angle,
                'width': width
            })
        
        # Sort roads by angle around intersection
        road_edges.sort(key=lambda x: x['angle'])
        
        # Connect adjacent roads with fillets (using intersection Z)
        for i in range(len(road_edges)):
            r1 = road_edges[i]
            r2 = road_edges[(i + 1) % len(road_edges)]
            
            radius = (r1['width'] + r2['width']) / 4
            builder.add_fillet_arc(center, r1['right_pt'], r2['left_pt'], radius, intersection_z)
    
    return builder.vertices, builder.faces

def get_utm_zone(lat: float, lon: float) -> int:
    """Determine UTM zone from latitude and longitude."""
    zone = int((lon + 180) / 6) + 1
    return zone

def get_utm_epsg(lat: float, lon: float) -> int:
    """Get EPSG code for UTM zone based on latitude and longitude."""
    zone = get_utm_zone(lat, lon)
    # Northern hemisphere: 32600 + zone, Southern hemisphere: 32700 + zone
    if lat >= 0:
        return 32600 + zone
    else:
        return 32700 + zone

def fetch_elevation_grid(south: float, north: float, west: float, east: float) -> Tuple[Optional[np.ndarray], float, float, float, float]:
    """
    Fetch elevation data from Mapbox Terrain-RGB for the given bbox.
    Returns: (elevation_grid, min_lat, min_lon, lat_res, lon_res)
    """
    # Use zoom level 14 for a good balance of resolution and tile count
    ZOOM = 14
    
    # Get the tiles covering the bbox
    tiles = list(mercantile.tiles(west, south, east, north, ZOOM))
    
    if not tiles:
        return None, 0, 0, 0, 0

    # Determine the bounds of the tile grid
    min_x = min(t.x for t in tiles)
    max_x = max(t.x for t in tiles)
    min_y = min(t.y for t in tiles)
    max_y = max(t.y for t in tiles)
    
    width_tiles = max_x - min_x + 1
    height_tiles = max_y - min_y + 1
    
    # Create a composite image
    tile_size = 256
    composite = Image.new('RGB', (width_tiles * tile_size, height_tiles * tile_size))
    
    for tile in tiles:
        url = f"https://api.mapbox.com/v4/mapbox.terrain-rgb/{tile.z}/{tile.x}/{tile.y}.pngraw?access_token={MAPBOX_TOKEN}"
        try:
            response = requests.get(url, stream=True)
            if response.status_code == 200:
                img = Image.open(BytesIO(response.content))
                # Paste into composite
                x_offset = (tile.x - min_x) * tile_size
                y_offset = (tile.y - min_y) * tile_size
                composite.paste(img, (x_offset, y_offset))
        except Exception as e:
            print(f"Error fetching tile {tile}: {e}")
            return None, 0, 0, 0, 0

    # Convert to numpy array
    data = np.array(composite)
    
    # Decode height: height = -10000 + ((R * 256 * 256 + G * 256 + B) * 0.1)
    r = data[:, :, 0].astype(np.float32)
    g = data[:, :, 1].astype(np.float32)
    b = data[:, :, 2].astype(np.float32)
    elevation = -10000 + ((r * 256 * 256 + g * 256 + b) * 0.1)
    
    # Calculate bounds of the composite image
    # Top-left of the top-left tile
    ul_tile = mercantile.Tile(min_x, min_y, ZOOM)
    ul_bounds = mercantile.bounds(ul_tile)
    max_lat_grid = ul_bounds.north
    min_lon_grid = ul_bounds.west
    
    # Bottom-right of the bottom-right tile
    lr_tile = mercantile.Tile(max_x, max_y, ZOOM)
    lr_bounds = mercantile.bounds(lr_tile)
    min_lat_grid = lr_bounds.south
    max_lon_grid = lr_bounds.east
    
    # Calculate resolution (degrees per pixel)
    lat_span = max_lat_grid - min_lat_grid
    lon_span = max_lon_grid - min_lon_grid
    
    pixel_height, pixel_width = elevation.shape
    lat_res = lat_span / pixel_height
    lon_res = lon_span / pixel_width
    
    # Crop to the requested bbox
    # Find indices
    start_row = int((max_lat_grid - north) / lat_res)
    end_row = int((max_lat_grid - south) / lat_res)
    start_col = int((west - min_lon_grid) / lon_res)
    end_col = int((east - min_lon_grid) / lon_res)
    
    # Clamp indices
    start_row = max(0, start_row)
    end_row = min(pixel_height, end_row)
    start_col = max(0, start_col)
    end_col = min(pixel_width, end_col)
    
    cropped_elevation = elevation[start_row:end_row, start_col:end_col]
    
    # Recalculate the precise bounds of the cropped area
    cropped_max_lat = max_lat_grid - start_row * lat_res
    cropped_min_lon = min_lon_grid + start_col * lon_res
    
    return cropped_elevation, cropped_max_lat, cropped_min_lon, lat_res, lon_res

def get_elevation_at_point(lat: float, lon: float, elevation_data: Tuple) -> float:
    """Sample elevation from the grid at a specific point. Clamps to edge if out of bounds."""
    grid, max_lat, min_lon, lat_res, lon_res = elevation_data
    
    if grid is None:
        return 0.0
        
    row = int((max_lat - lat) / lat_res)
    col = int((lon - min_lon) / lon_res)
    
    rows, cols = grid.shape
    
    # Clamp to valid range to avoid dropping to zero at edges
    row = max(0, min(row, rows - 1))
    col = max(0, min(col, cols - 1))
    
    return float(grid[row, col])

def generate_obj_file(center_lat: float, center_lon: float, square_size_km: float, layers: List[str] = ["Buildings", "Roads"], 
                      contour_interval: float = 0, use_utm_projection: bool = False) -> Tuple[Optional[str], Optional[str], str]:
    """Generate OBJ file with buildings, roads (as polygon footprints), and topography.
    
    Args:
        center_lat: Latitude of center
        center_lon: Longitude of center
        square_size_km: Size of square area in km
        layers: List of layers to include
        contour_interval: Elevation interval for contour lines in meters (0 = disabled)
        use_utm_projection: If True, reproject all geometry to UTM
    """
    
    include_buildings = "Buildings" in layers
    include_roads = "Roads" in layers
    include_terrain = "Terrain" in layers
    
    # Determine UTM EPSG if reprojection is requested
    utm_epsg = None
    if use_utm_projection:
        utm_epsg = get_utm_epsg(center_lat, center_lon)
        utm_zone = get_utm_zone(center_lat, center_lon)
        status = f"UTM Projection Enabled: Zone {utm_zone} (EPSG:{utm_epsg})\n"
    else:
        status = ""
    
    # Validate inputs
    if not (-90 <= center_lat <= 90):
        return None, None, "Error: Latitude must be between -90 and 90"
    if not (-180 <= center_lon <= 180):
        return None, None, "Error: Longitude must be between -180 and 180"
    if not (0.5 <= square_size_km <= 5):
        return None, None, "Error: Square size must be between 0.5 and 5 km"
    
    status += f"Fetching OSM data for {center_lat}, {center_lon} ({square_size_km}km)...\n"
    half_size_km = square_size_km / 2
    lat_offset, lon_offset = km_to_degrees(half_size_km, center_lat)
    south = center_lat - lat_offset
    north = center_lat + lat_offset
    west = center_lon - lon_offset
    east = center_lon + lon_offset
    
    # Fetch Elevation Data (needed for terrain OR for placing objects on terrain)
    elevation_data = None
    if include_terrain or include_buildings or include_roads:
        status += "Fetching elevation data...\n"
        elevation_grid, max_lat_grid, min_lon_grid, lat_res, lon_res = fetch_elevation_grid(south, north, west, east)
        if elevation_grid is not None:
            elevation_data = (elevation_grid, max_lat_grid, min_lon_grid, lat_res, lon_res)
            status += f"Received elevation grid: {elevation_grid.shape}\n"
        else:
            status += "Warning: Failed to fetch elevation data. Using flat terrain.\n"

    # Fetch OSM Data
    data = None
    if include_buildings or include_roads:
        data = fetch_osm_data(center_lat, center_lon, square_size_km)
        if not data:
            status += "Failed to fetch OSM data (or none found)\n"
        else:
            status += f"Received {len(data['elements'])} OSM elements\n"
    
    # Parse data
    nodes = {}
    buildings = []
    roads = []
    
    if data:
        for element in data['elements']:
            if element['type'] == 'node':
                nodes[element['id']] = (element['lat'], element['lon'])
            elif element['type'] == 'way' and 'tags' in element:
                if include_buildings and 'building' in element['tags'] and 'nodes' in element:
                    height = get_building_height(element['tags'])
                    buildings.append({'nodes': element['nodes'], 'height': height})
                elif include_roads and 'highway' in element['tags'] and 'nodes' in element:
                    roads.append({'nodes': element['nodes'], 'type': element['tags']['highway']})
    
    status += f"Found {len(buildings)} buildings\nFound {len(roads)} road centerlines\n"
    
    if not buildings and not roads and not include_terrain:
        return None, None, status + "Nothing to generate"
    
    # Create transformer once for UTM projection
    transformer_to_utm = None
    if use_utm_projection:
        transformer_to_utm = Transformer.from_crs("EPSG:4326", f"EPSG:{utm_epsg}", always_xy=True)
    
    # Prepare roads data for smooth footprint generation
    roads_with_points = []
    if include_roads and roads:
        for road in roads:
            road_type = road.get('type', 'residential')
            width = get_road_width(road_type)
            
            points = []  # Local coordinates (x, y)
            lat_lon_points = []  # Original lat/lon for elevation lookup
            
            for node_id in road['nodes']:
                if node_id in nodes:
                    lat, lon = nodes[node_id]
                    lat_lon_points.append((lat, lon))
                    
                    if use_utm_projection:
                        utm_x, utm_y = transformer_to_utm.transform(lon, lat)
                        # Store raw UTM coords, will center later
                        points.append((utm_x, utm_y))
                    else:
                        x, y = latlon_to_meters(lat, lon, center_lat, center_lon)
                        points.append((x, y))
            
            if len(points) >= 2:
                roads_with_points.append({
                    'points': points,
                    'width': width,
                    'type': road_type,
                    'lat_lon_points': lat_lon_points
                })
    
    # Generate files in a local directory for proper downloads
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    base_name = f"area_{square_size_km}km_{timestamp}"
    
    # Use temp directory - Gradio will handle cleanup
    #output_dir = tempfile.gettempdir()
    
    # Create output directory if it doesn't exist
    output_dir = "generated_models"
    os.makedirs(output_dir, exist_ok=True)
    
    obj_file = os.path.join(output_dir, f"{base_name}.obj")
    mtl_file = os.path.join(output_dir, f"{base_name}.mtl")
    
    status += "Generating 3D model...\n"
    
    # Write OBJ file
    mtl_filename = os.path.basename(mtl_file)  # Use only filename, not full path
    with open(obj_file, 'w') as f:
        f.write(f"# OSM 3D Export - {square_size_km}km x {square_size_km}km\n")
        f.write(f"# Center: {center_lat}, {center_lon}\n")
        f.write(f"mtllib {mtl_filename}\n\n")
        
        # Normals
        f.write("vn 0.0 0.0 -1.0\n")  # Bottom
        f.write("vn 0.0 0.0 1.0\n")   # Top
        f.write("\n")
        
        vertex_count = 0
        
        # Calculate center point in UTM for offset (to keep coordinates manageable)
        center_utm_x, center_utm_y = 0, 0
        if use_utm_projection:
            # Reuse transformer if already created, otherwise create it
            if transformer_to_utm is None:
                transformer_to_utm = Transformer.from_crs("EPSG:4326", f"EPSG:{utm_epsg}", always_xy=True)
            center_utm_x, center_utm_y = transformer_to_utm.transform(center_lon, center_lat)
        
        # --- TERRAIN GENERATION ---
        if include_terrain and elevation_data:
            f.write("g Terrain\nusemtl Ground\n")
            grid, max_lat, min_lon, lat_res, lon_res = elevation_data
            rows, cols = grid.shape
            
            # Use step 1 for maximum quality/detail
            step = 1
            
            # Create vertices
            # We need to map grid indices back to meters relative to center
            grid_vertices = {} # (row, col) -> vertex_index
            
            for r in range(0, rows, step):
                for c in range(0, cols, step):
                    lat = max_lat - r * lat_res
                    lon = min_lon + c * lon_res
                    
                    if use_utm_projection:
                        # Convert lat/lon to UTM and center
                        utm_x, utm_y = transformer_to_utm.transform(lon, lat)
                        x = utm_x - center_utm_x
                        y = utm_y - center_utm_y
                    else:
                        # Convert lat/lon to local meters
                        x, y = latlon_to_meters(lat, lon, center_lat, center_lon)
                    
                    z = float(grid[r, c])
                    
                    f.write(f"v {x:.2f} {y:.2f} {z:.2f}\n")
                    vertex_count += 1
                    grid_vertices[(r, c)] = vertex_count
            
            # Create faces
            for r in range(0, rows - step, step):
                for c in range(0, cols - step, step):
                    # v1 v2
                    # v3 v4
                    v1 = grid_vertices.get((r, c))
                    v2 = grid_vertices.get((r, c + step))
                    v3 = grid_vertices.get((r + step, c))
                    v4 = grid_vertices.get((r + step, c + step))
                    
                    if v1 and v2 and v3 and v4:
                        # Two triangles
                        f.write(f"f {v1} {v3} {v2}\n")
                        f.write(f"f {v2} {v3} {v4}\n")
            f.write("\n")
            
            # --- CONTOUR LINES ---
            if contour_interval > 0:
                f.write("g Contours\nusemtl Contour\n")
                
                min_elev = float(grid.min())
                max_elev = float(grid.max())
                
                # Generate contour lines at each interval
                contour_levels = np.arange(
                    np.floor(min_elev / contour_interval) * contour_interval,
                    np.ceil(max_elev / contour_interval) * contour_interval + contour_interval,
                    contour_interval
                )
                
                # Track first point of each contour level for label placement
                contour_label_points = {}
                
                for level in contour_levels:
                    level_points = []
                    # Simple contour extraction using marching squares concept
                    # For each grid cell, check if contour passes through
                    for r in range(0, rows - 1, 1):  # Use step 1 for accurate contours
                        for c in range(0, cols - 1, 1):
                            # Get the 4 corners of this cell
                            z1 = grid[r, c]
                            z2 = grid[r, c + 1]
                            z3 = grid[r + 1, c]
                            z4 = grid[r + 1, c + 1]
                            
                            # Check if contour crosses this cell
                            zmin = min(z1, z2, z3, z4)
                            zmax = max(z1, z2, z3, z4)
                            
                            if zmin <= level <= zmax:
                                # Contour passes through - create a line segment
                                # Simple approach: if level is between corners, interpolate
                                points = []
                                
                                # Check each edge
                                # Top edge (r, c) to (r, c+1)
                                if (z1 <= level <= z2) or (z2 <= level <= z1):
                                    if z1 != z2:
                                        t = (level - z1) / (z2 - z1)
                                        lat = max_lat - r * lat_res
                                        lon = min_lon + (c + t) * lon_res
                                        if use_utm_projection:
                                            utm_x, utm_y = transformer_to_utm.transform(lon, lat)
                                            x = utm_x - center_utm_x
                                            y = utm_y - center_utm_y
                                        else:
                                            x, y = latlon_to_meters(lat, lon, center_lat, center_lon)
                                        points.append((x, y, level))
                                
                                # Right edge (r, c+1) to (r+1, c+1)
                                if (z2 <= level <= z4) or (z4 <= level <= z2):
                                    if z2 != z4:
                                        t = (level - z2) / (z4 - z2)
                                        lat = max_lat - (r + t) * lat_res
                                        lon = min_lon + (c + 1) * lon_res
                                        if use_utm_projection:
                                            utm_x, utm_y = transformer_to_utm.transform(lon, lat)
                                            x = utm_x - center_utm_x
                                            y = utm_y - center_utm_y
                                        else:
                                            x, y = latlon_to_meters(lat, lon, center_lat, center_lon)
                                        points.append((x, y, level))
                                
                                # Bottom edge (r+1, c) to (r+1, c+1)
                                if (z3 <= level <= z4) or (z4 <= level <= z3):
                                    if z3 != z4:
                                        t = (level - z3) / (z4 - z3)
                                        lat = max_lat - (r + 1) * lat_res
                                        lon = min_lon + (c + t) * lon_res
                                        if use_utm_projection:
                                            utm_x, utm_y = transformer_to_utm.transform(lon, lat)
                                            x = utm_x - center_utm_x
                                            y = utm_y - center_utm_y
                                        else:
                                            x, y = latlon_to_meters(lat, lon, center_lat, center_lon)
                                        points.append((x, y, level))
                                
                                # Left edge (r, c) to (r+1, c)
                                if (z1 <= level <= z3) or (z3 <= level <= z1):
                                    if z1 != z3:
                                        t = (level - z1) / (z3 - z1)
                                        lat = max_lat - (r + t) * lat_res
                                        lon = min_lon + c * lon_res
                                        if use_utm_projection:
                                            utm_x, utm_y = transformer_to_utm.transform(lon, lat)
                                            x = utm_x - center_utm_x
                                            y = utm_y - center_utm_y
                                        else:
                                            x, y = latlon_to_meters(lat, lon, center_lat, center_lon)
                                        points.append((x, y, level))
                                
                                # Create line segments from points (take first 2)
                                if len(points) >= 2:
                                    for i in range(0, len(points) - 1, 2):
                                        if i + 1 < len(points):
                                            p1 = points[i]
                                            p2 = points[i + 1]
                                            
                                            # Store first point for label
                                            if level not in contour_label_points:
                                                contour_label_points[level] = p1
                                            
                                            # Add slight offset above terrain
                                            z_offset = level + 0.5
                                            
                                            v1 = vertex_count + 1
                                            f.write(f"v {p1[0]:.2f} {p1[1]:.2f} {z_offset:.2f}\n")
                                            vertex_count += 1
                                            v2 = vertex_count + 1
                                            f.write(f"v {p2[0]:.2f} {p2[1]:.2f} {z_offset:.2f}\n")
                                            vertex_count += 1
                                            
                                            # Line as degenerate face
                                            f.write(f"l {v1} {v2}\n")
                
                # Add elevation labels as 3D text (using simple line-based text)
                f.write("\ng ContourLabels\nusemtl Contour\n")
                for level, (x, y, z) in contour_label_points.items():
                    # Create simple text label using comment (visible in some viewers)
                    # Also create a small marker point
                    label_text = f"{int(level)}m"
                    f.write(f"# Elevation: {label_text} at ({x:.1f}, {y:.1f})\n")
                    
                    # Create a small vertical line as a marker
                    marker_height = 2.0
                    v1 = vertex_count + 1
                    f.write(f"v {x:.2f} {y:.2f} {z + 0.5:.2f}\n")
                    vertex_count += 1
                    v2 = vertex_count + 1
                    f.write(f"v {x:.2f} {y:.2f} {z + marker_height:.2f}\n")
                    vertex_count += 1
                    f.write(f"l {v1} {v2}\n")
                    
                f.write("\n")

        # --- BUILDINGS GENERATION ---
        if include_buildings:
            for i, building in enumerate(buildings, 1):
                coords = []
                base_z = 0.0
                
                # Calculate base height (only if terrain is enabled)
                node_elevations = []
                for node_id in building['nodes']:
                    if node_id in nodes:
                        lat, lon = nodes[node_id]
                        if use_utm_projection:
                            # Reproject to UTM and center around the center point
                            utm_x, utm_y = transformer_to_utm.transform(lon, lat)
                            x = utm_x - center_utm_x
                            y = utm_y - center_utm_y
                        else:
                            x, y = latlon_to_meters(lat, lon, center_lat, center_lon)
                        coords.append((x, y))
                        # Only collect elevations if terrain is enabled
                        if include_terrain and elevation_data:
                            node_elevations.append(get_elevation_at_point(lat, lon, elevation_data))
                
                if not coords:
                    continue
                    
                # Set base Z to the minimum elevation of the footprint to avoid floating
                # Only if terrain is enabled, otherwise base_z stays at 0.0
                if include_terrain and node_elevations:
                    base_z = min(node_elevations)
                
                if len(coords) < 3:
                    continue
                if coords[0] == coords[-1]:
                    coords = coords[:-1]
                if len(coords) < 3:
                    continue
                
                height = building['height']
                top_z = base_z + height
                
                f.write(f"g Buildings\no Building_{i}\nusemtl White\n")
                
                # Bottom vertices
                bottom_start = vertex_count + 1
                for x, y in coords:
                    f.write(f"v {x:.2f} {y:.2f} {base_z:.2f}\n")
                    vertex_count += 1
                
                # Top vertices
                top_start = vertex_count + 1
                for x, y in coords:
                    f.write(f"v {x:.2f} {y:.2f} {top_z:.2f}\n")
                    vertex_count += 1
                
                num_points = len(coords)
                
                # Roof
                f.write("f")
                for j in range(num_points):
                    f.write(f" {top_start + j}//2")
                f.write("\n")
                
                # Walls
                for j in range(num_points):
                    next_j = (j + 1) % num_points
                    v1, v2 = bottom_start + j, bottom_start + next_j
                    v3, v4 = top_start + next_j, top_start + j
                    # Simple normals (using 1 for side for now to simplify)
                    f.write(f"f {v1} {v2} {v3} {v4}\n")
                f.write("\n")
        
        # --- ROADS GENERATION (Smooth Footprints with V2 Algorithm) ---
        if include_roads and roads_with_points:
            status += "Generating smooth road footprints...\n"
            
            # Apply UTM centering offset to road points if using UTM projection
            if use_utm_projection:
                for road in roads_with_points:
                    road['points'] = [(x - center_utm_x, y - center_utm_y) for x, y in road['points']]
            
            # Generate smooth road geometry
            # Only use terrain elevation if terrain layer is enabled
            road_vertices, road_faces = generate_smooth_roads(
                roads_with_points, 
                elevation_data, 
                get_elevation_at_point,
                use_terrain=include_terrain
            )
            
            if road_vertices and road_faces:
                f.write("g Roads\nusemtl Road\n")
                
                # Write road vertices
                road_vertex_start = vertex_count + 1
                for x, y, z in road_vertices:
                    f.write(f"v {x:.2f} {y:.2f} {z:.2f}\n")
                    vertex_count += 1
                
                # Write road faces (indices need to be offset)
                for i0, i1, i2 in road_faces:
                    # road_faces uses 1-based indexing from builder, add our offset
                    f.write(f"f {road_vertex_start + i0 - 1} {road_vertex_start + i1 - 1} {road_vertex_start + i2 - 1}\n")
                
                f.write("\n")
                status += f"Generated {len(road_vertices)} road vertices, {len(road_faces)} faces\n"
            else:
                status += "Warning: No road geometry generated\n"
    
    # Write MTL file
    with open(mtl_file, 'w') as f:
        f.write("# Materials\n\n")
        # Off-white for Buildings
        f.write("newmtl White\nKa 0.95 0.95 0.9\nKd 0.98 0.98 0.95\n")
        f.write("Ks 0.1 0.1 0.1\nNs 10.0\nd 1.0\nillum 2\n\n")
        # Dark gray for Roads
        f.write("newmtl Road\nKa 0.15 0.15 0.15\nKd 0.25 0.25 0.25\n")
        f.write("Ks 0.05 0.05 0.05\nNs 5.0\nd 1.0\nillum 2\n\n")
        # Light gray for Ground/Terrain
        f.write("newmtl Ground\nKa 0.7 0.7 0.7\nKd 0.75 0.75 0.75\n")
        f.write("Ks 0.0 0.0 0.0\nNs 0.0\nd 1.0\nillum 1\n\n")
        # Black for Contour lines
        f.write("newmtl Contour\nKa 0.0 0.0 0.0\nKd 0.05 0.05 0.05\n")
        f.write("Ks 0.0 0.0 0.0\nNs 0.0\nd 1.0\nillum 1\n")
    
    status += f"Generated {len(buildings)} buildings, {len(roads)} roads\n"
    status += f"Files: {obj_file}, {mtl_file}"
    
    return obj_file, mtl_file, status

# Gradio Interface with Map and 3D Viewer
with gr.Blocks(title="OSM 3D Generator", theme=gr.themes.Soft(), css="""
    #map-container { height: 400px; }
    .model-3d { height: 500px; }
    #map-image .image-container { border: none !important; }
    #map-image .upload-container { display: none !important; }
    button[aria-label*="Clear"] { display: none !important; }
    button[aria-label*="Download"] { display: none !important; }
""") as app:
    gr.Markdown("""
    # OSM 3D Environment Generator
    Generate 3D building models from OpenStreetMap data
    """)
    
    with gr.Row():
        # Left Column - Map and Controls
        with gr.Column(scale=1):
            gr.Markdown("### Location Selector")
            
            # Static map preview using Mapbox Static Images API with bounding box
            def get_map_url(lat, lon, square_size_km):
                """Generate Mapbox static image URL with bounding box overlay."""
                # Calculate bounding box coordinates
                lat_offset, lon_offset = km_to_degrees(square_size_km / 2, lat)
                
                # Dynamic zoom level based on area size to ensure bbox fits in view
                if square_size_km <= 1.0:
                    zoom = 14
                elif square_size_km <= 2.0:
                    zoom = 13
                elif square_size_km <= 3.0:
                    zoom = 12
                elif square_size_km <= 4.0:
                    zoom = 12
                else:  # 5.0 km
                    zoom = 11.5
                
                # Create GeoJSON overlay for the bounding box
                geojson = {
                    "type": "FeatureCollection",
                    "features": [{
                        "type": "Feature",
                        "geometry": {
                            "type": "Polygon",
                            "coordinates": [[
                                [lon - lon_offset, lat - lat_offset],
                                [lon + lon_offset, lat - lat_offset],
                                [lon + lon_offset, lat + lat_offset],
                                [lon - lon_offset, lat + lat_offset],
                                [lon - lon_offset, lat - lat_offset]
                            ]]
                        },
                        "properties": {
                            "stroke": "#ff0000",
                            "stroke-width": 3,
                            "fill": "#ff0000",
                            "fill-opacity": 0.1
                        }
                    }]
                }
                
                import json
                import urllib.parse
                geojson_str = json.dumps(geojson)
                geojson_encoded = urllib.parse.quote(geojson_str)
                
                # Use satellite-v9 for pure satellite imagery without labels
                return f"https://api.mapbox.com/styles/v1/mapbox/satellite-v9/static/geojson({geojson_encoded}),pin-s+ff0000({lon},{lat})/{lon},{lat},{zoom},0/600x400@2x?access_token={MAPBOX_TOKEN}"
            
            map_image = gr.Image(
                value=get_map_url(52.52309, 13.41657, 2.0),
                label="Map Preview",
                height=400,
                interactive=False,
                elem_id="map-image",
                show_label=False,
                show_download_button=False,
                show_share_button=False
            )
            
            gr.Markdown("""
            **Select Location:**  
            Enter coordinates manually or use [Google Maps](https://maps.google.com) / [OpenStreetMap](https://www.openstreetmap.org) to find coordinates.  
            Red box shows the fetch area.
            """)
            
            gr.Markdown("### Coordinates")
            with gr.Row():
                lat = gr.Number(label="Latitude", value=52.52309, precision=6)
                lon = gr.Number(label="Longitude", value=13.41657, precision=6)
            
            update_map_btn = gr.Button("Update Map Preview", size="sm")
            
            size = gr.Slider(label="Area Size (km)", minimum=0.5, maximum=5, value=2.0, step=0.5)
            
            layers = gr.CheckboxGroup(
                choices=["Buildings", "Roads", "Terrain"],
                value=["Buildings", "Roads"],
                label="Layers to Generate"
            )
            
            with gr.Accordion("Terrain Options", open=False):
                contour_interval = gr.Slider(
                    label="Contour Line Interval (meters)",
                    minimum=0,
                    maximum=50,
                    value=0,
                    step=1,
                    info="0 = disabled, 2-10m for detailed, 10-50m for overview. Labels show elevation."
                )
            
            with gr.Accordion("Export Options", open=False):
                use_utm = gr.Checkbox(
                    label="Reproject to UTM (for CAD/GIS alignment)",
                    value=False,
                    info="Enables coordinate system matching with UTM-projected satellite imagery"
                )
            
            generate_btn = gr.Button("Generate 3D Model", variant="primary", size="lg")
            
            status_output = gr.Textbox(label="Status", lines=8, interactive=False)
        
        # Right Column - 3D Viewer and Downloads
        with gr.Column(scale=1):
            gr.Markdown("### 3D Model Viewer")
            
            model_viewer = gr.Model3D(
                label="3D Preview",
                height=500,
                interactive=False
            )
            
            gr.Markdown("### Download Files")
            with gr.Row():
                obj_output = gr.File(label="OBJ File", file_count="single")
                mtl_output = gr.File(label="MTL File", file_count="single")
    
    with gr.Row():
        gr.Markdown("""
        ### How to Use
        1. Enter coordinates or use Google Maps/OpenStreetMap to find a location
        2. Click "Update Map Preview" to visualize the selected area
        3. Adjust the area size using the slider (automatically updates map)
        4. Select which layers to include (Buildings, Roads, Terrain)
        5. Click "Generate 3D Model" to create buildings and roads
        6. Preview the model in 3D and download OBJ/MTL files
        
        ### Features
        - Pure satellite imagery without labels
        - Red bounding box shows exact fetch area
        - Real-time 3D preview in browser
        - Building heights from OSM data
        - **Road footprints** with realistic width (buffered polygons for CAD/laser-etch)
        - **High-precision Terrain** from Mapbox Terrain-RGB
        - Optional UTM reprojection for GIS/CAD alignment
        - Professional materials and normals
        - Compatible with Rhino, Blender, 3ds Max
        """)
    
    # Quick examples
    with gr.Row():
        gr.Markdown("### Quick Examples")
        examples_data = [
            [52.52309, 13.41657, 2.0],   # Berlin
            [40.74882, -73.98543, 1.0],  # NYC
            [48.85884, 2.29435, 1.5],    # Paris
            [51.50073, -0.12463, 2.0],   # London
            [35.67610, 139.65031, 1.5],  # Tokyo
        ]
        gr.Examples(
            examples=examples_data,
            inputs=[lat, lon, size],
            label="Click to load example locations"
        )
    
    # Event handlers
    def handle_generation(lat_val, lon_val, size_val, layers_val, contour_interval_val, utm_val):
        """Generate model and return files + preview."""
        obj_file, mtl_file, status = generate_obj_file(
            lat_val, lon_val, size_val, layers_val,
            contour_interval=float(contour_interval_val),
            use_utm_projection=bool(utm_val)
        )
        
        if obj_file:
            return obj_file, mtl_file, obj_file, status
        else:
            return None, None, None, status
    
    def update_map_preview(lat_val, lon_val, size_val):
        """Update map image when coordinates or size change."""
        return get_map_url(lat_val, lon_val, size_val)
    
    # Connect button events
    generate_btn.click(
        fn=handle_generation,
        inputs=[lat, lon, size, layers, contour_interval, use_utm],
        outputs=[obj_output, mtl_output, model_viewer, status_output]
    )
    
    update_map_btn.click(
        fn=update_map_preview,
        inputs=[lat, lon, size],
        outputs=[map_image]
    )
    
    # Auto-update map when size slider changes
    size.change(
        fn=update_map_preview,
        inputs=[lat, lon, size],
        outputs=[map_image]
    )

if __name__ == "__main__":
    app.launch()