plot_clusters.py

import xml.etree.ElementTree as ET
import plotly.graph_objects as go
import gzip
import math
import sys
import os
import argparse as ap

# Role-based color mapping with 50% opacity
ROLE_COLORS = {
    'STOP_POSITION': 'rgba(255, 0, 0, 0.5)',        # red
    'PLATFORM_EDGE': 'rgba(255, 165, 0, 0.5)',      # orange
    'PLATFORM':      'rgba(0, 180, 0, 0.5)',         # green
    'STATION':       'rgba(128, 0, 128, 0.5)',       # purple
    'STOP_AREA':     'rgba(255, 0, 255, 0.5)',       # magenta
    'TRACK':         'rgba(0, 0, 0, 0.5)',           # black
    'OUTER':         'rgba(139, 69, 19, 0.5)',       # brown
    'UNKNOWN':       'rgba(30, 120, 180, 0.5)',      # blue
    'NOT_EVALUATED': 'rgba(128, 128, 128, 0.5)',     # grey
}

def get_role_color(role_str):
    return ROLE_COLORS.get(role_str, ROLE_COLORS['UNKNOWN'])

# Cluster boundary colors — one distinct color per cluster (cycled)
BOUNDARY_COLORS = [
    'rgba(255, 0, 0, 0.2)',
    'rgba(0, 128, 0, 0.2)',
    'rgba(0, 0, 255, 0.2)',
    'rgba(255, 165, 0, 0.2)',
    'rgba(128, 0, 128, 0.2)',
    'rgba(0, 128, 128, 0.2)',
    'rgba(255, 0, 255, 0.2)',
    'rgba(139, 69, 19, 0.2)',
    'rgba(0, 0, 139, 0.2)',
    'rgba(184, 134, 11, 0.2)',
]


def convex_hull(points):
    """Compute convex hull of 2D points using Andrew's monotone chain."""
    points = sorted(set(points))
    if len(points) <= 1:
        return points
    # Build lower hull
    lower = []
    for p in points:
        while len(lower) >= 2 and cross(lower[-2], lower[-1], p) <= 0:
            lower.pop()
        lower.append(p)
    # Build upper hull
    upper = []
    for p in reversed(points):
        while len(upper) >= 2 and cross(upper[-2], upper[-1], p) <= 0:
            upper.pop()
        upper.append(p)
    return lower[:-1] + upper[:-1]


def cross(o, a, b):
    return (a[0] - o[0]) * (b[1] - o[1]) - (a[1] - o[1]) * (b[0] - o[0])


def buffer_hull(hull_points, radius_m, center_lat):
    """Buffer a convex hull by radius_m meters, approximated in lat/lon degrees."""
    if not hull_points:
        return []
    # Convert meters to degrees
    lat_per_m = 1.0 / 111320.0
    lon_per_m = 1.0 / (111320.0 * math.cos(math.radians(center_lat)))
    dlat = radius_m * lat_per_m
    dlon = radius_m * lon_per_m

    if len(hull_points) == 1:
        # Single point: create a circle
        cx, cy = hull_points[0]
        return _circle(cx, cy, dlat, dlon, 36)

    # Expand each edge outward and round corners
    n = len(hull_points)
    buffered = []
    for i in range(n):
        p0 = hull_points[i]
        p1 = hull_points[(i + 1) % n]
        # Direction vector
        dx = p1[0] - p0[0]
        dy = p1[1] - p0[1]
        length = math.sqrt(dx * dx + dy * dy)
        if length == 0:
            continue
        # Outward normal (counterclockwise hull → normal points outward)
        nx = -dy / length
        ny = dx / length
        # Offset points
        off_x = nx * dlat
        off_y = ny * dlon

        # Add rounded corner at p0
        prev = hull_points[(i - 1) % n]
        dx_prev = p0[0] - prev[0]
        dy_prev = p0[1] - prev[1]
        len_prev = math.sqrt(dx_prev * dx_prev + dy_prev * dy_prev)
        if len_prev > 0:
            nx_prev = -dy_prev / len_prev
            ny_prev = dx_prev / len_prev
            angle_prev = math.atan2(ny_prev * dlon, nx_prev * dlat)
            angle_curr = math.atan2(off_y, off_x)
            # Generate arc points for rounded corner
            if angle_prev > angle_curr:
                angle_curr += 2 * math.pi
            steps = max(2, int((angle_curr - angle_prev) / (math.pi / 8)))
            for s in range(steps + 1):
                a = angle_prev + (angle_curr - angle_prev) * s / steps
                bx = p0[0] + dlat * math.cos(a)
                by = p0[1] + dlon * math.sin(a)
                buffered.append((bx, by))
        else:
            buffered.append((p0[0] + off_x, p0[1] + off_y))
        buffered.append((p1[0] + off_x, p1[1] + off_y))

    if not buffered:
        # Fallback: offset each hull point outward from centroid
        cx = sum(p[0] for p in hull_points) / n
        cy = sum(p[1] for p in hull_points) / n
        for p in hull_points:
            dx = p[0] - cx
            dy = p[1] - cy
            length = math.sqrt(dx * dx + dy * dy)
            if length > 0:
                buffered.append((p[0] + dlat * dx / length, p[1] + dlon * dy / length))
            else:
                buffered.append(p)

    # Re-compute convex hull of buffered points for a clean shape
    return convex_hull(buffered)


def _circle(cx, cy, rx, ry, n=36):
    """Generate circle points."""
    pts = []
    for i in range(n):
        angle = 2 * math.pi * i / n
        pts.append((cx + rx * math.cos(angle), cy + ry * math.sin(angle)))
    return pts


def plot_clusters(xml_file, output_path=None):
    try:
        if xml_file.endswith('.gz'):
            with gzip.open(xml_file, 'rb') as fh:
                tree = ET.parse(fh)
        else:
            tree = ET.parse(xml_file)
        root = tree.getroot()
    except FileNotFoundError:
        print(f"Error: File {xml_file} not found.")
        return
    except ET.ParseError as e:
        print(f"Error: Failed to parse {xml_file}. Details: {e}")
        return

    fig = go.Figure()

    clusters = root.findall('cluster')
    print(f"Found {len(clusters)} clusters in XML.")

    # Group nodes by role
    nodes_by_role = {}
    # Group way lines by role
    way_lines_by_role = {}
    # Group relation lines by role
    rel_lines_by_role = {}
    # GTFS stops (always yellow)
    gtfs_lats, gtfs_lons, gtfs_texts = [], [], []

    all_lats = []
    all_lons = []

    # Per-cluster points for boundary shapes
    cluster_points = {}  # cluster_id -> [(lat, lon), ...]

    for i, cluster in enumerate(clusters):
        cluster_id = cluster.get('id')
        cluster_pts = []

        # Lookup maps: member ID -> its own role (used when plotting relation members)
        node_roles = {}  # node_id_str -> role_str
        way_roles = {}   # way_id_str -> role_str

        # Collect OSM Nodes — skip REJECT status
        for node in cluster.findall('osm_node'):
            status = node.get('status', 'UNKNOWN')
            if status == 'REJECT':
                continue
            lat = float(node.get('lat'))
            lon = float(node.get('lon'))
            role = node.get('role', 'UNKNOWN')

            if role not in nodes_by_role:
                nodes_by_role[role] = {'lats': [], 'lons': [], 'texts': []}
            nodes_by_role[role]['lats'].append(lat)
            nodes_by_role[role]['lons'].append(lon)
            nodes_by_role[role]['texts'].append(
                f"Cluster {cluster_id}<br>Node {node.get('id')}<br>Role: {role}<br>Status: {status}"
            )
            all_lats.append(lat)
            all_lons.append(lon)
            cluster_pts.append((lat, lon))
            node_roles[node.get('id')] = role

        # Collect OSM Ways — skip REJECT status
        for way in cluster.findall('osm_way'):
            status = way.get('status', 'UNKNOWN')
            if status == 'REJECT':
                continue
            role = way.get('role', 'UNKNOWN')
            nodes_wrapper = way.find('nodes')
            if nodes_wrapper is None:
                continue
            nodes = nodes_wrapper.findall('node')
            if not nodes:
                continue

            if role not in way_lines_by_role:
                way_lines_by_role[role] = {'lats': [], 'lons': []}

            for node in nodes:
                try:
                    lat = float(node.get('lat'))
                    lon = float(node.get('lon'))
                    way_lines_by_role[role]['lats'].append(lat)
                    way_lines_by_role[role]['lons'].append(lon)
                    all_lats.append(lat)
                    all_lons.append(lon)
                    cluster_pts.append((lat, lon))
                except (ValueError, TypeError):
                    pass
            way_lines_by_role[role]['lats'].append(None)
            way_lines_by_role[role]['lons'].append(None)
            way_roles[way.get('id')] = role

        # Collect OSM Relations — skip REJECT status
        for relation in cluster.findall('osm_relation'):
            status = relation.get('status', 'UNKNOWN')
            if status == 'REJECT':
                continue
            rel_role = relation.get('role', 'UNKNOWN')

            # Node and way members — use each member's own role, not the relation's role
            for member in relation.findall('members/member'):
                m_type = member.get('type')
                m_ref = member.get('ref')

                if m_type == 'node':
                    details = member.find('details')
                    if details is not None:
                        lat_str = details.get('lat')
                        lon_str = details.get('lon')
                        if lat_str and lon_str:
                            lat = float(lat_str)
                            lon = float(lon_str)
                            # Use node's own role if known, otherwise fall back to relation role
                            member_role = node_roles.get(m_ref, rel_role)
                            if member_role not in nodes_by_role:
                                nodes_by_role[member_role] = {'lats': [], 'lons': [], 'texts': []}
                            nodes_by_role[member_role]['lats'].append(lat)
                            nodes_by_role[member_role]['lons'].append(lon)
                            nodes_by_role[member_role]['texts'].append(
                                f"Cluster {cluster_id}<br>Rel {relation.get('id')}<br>Node {m_ref}<br>Role: {member_role}"
                            )
                            all_lats.append(lat)
                            all_lons.append(lon)
                            cluster_pts.append((lat, lon))

                elif m_type == 'way':
                    details = member.find('details')
                    if details is not None:
                        nodes_wrapper = details.find('nodes')
                        if nodes_wrapper is not None:
                            nodes = nodes_wrapper.findall('node')
                            if not nodes:
                                continue

                            # Use way's own role if known, otherwise fall back to relation role
                            member_role = way_roles.get(m_ref, rel_role)
                            if member_role not in rel_lines_by_role:
                                rel_lines_by_role[member_role] = {'lats': [], 'lons': []}

                            for node in nodes:
                                try:
                                    lat = float(node.get('lat'))
                                    lon = float(node.get('lon'))
                                    rel_lines_by_role[member_role]['lats'].append(lat)
                                    rel_lines_by_role[member_role]['lons'].append(lon)
                                    all_lats.append(lat)
                                    all_lons.append(lon)
                                    cluster_pts.append((lat, lon))
                                except (ValueError, TypeError):
                                    pass
                            rel_lines_by_role[member_role]['lats'].append(None)
                            rel_lines_by_role[member_role]['lons'].append(None)

        # Collect GTFS Stops
        for stop in cluster.findall('gtfs_stop'):
            lat = float(stop.get('lat'))
            lon = float(stop.get('lon'))
            gtfs_lats.append(lat)
            gtfs_lons.append(lon)
            gtfs_texts.append(
                f"Cluster {cluster_id}<br>GTFS: {stop.get('name')} ({stop.get('id')})"
            )
            all_lats.append(lat)
            all_lons.append(lon)
            cluster_pts.append((lat, lon))

        if cluster_pts:
            cluster_points[cluster_id] = cluster_pts

    # --- Add cluster boundary shapes first (behind everything) ---
    for idx, (cluster_id, pts) in enumerate(cluster_points.items()):
        if len(pts) < 1:
            continue
        hull = convex_hull(pts)
        center_lat = sum(p[0] for p in pts) / len(pts)
        buffered = buffer_hull(hull, 100, center_lat)  # 100m radius

        if buffered:
            # Close the polygon
            boundary_lats = [p[0] for p in buffered] + [buffered[0][0]]
            boundary_lons = [p[1] for p in buffered] + [buffered[0][1]]
            color = BOUNDARY_COLORS[idx % len(BOUNDARY_COLORS)]
            fig.add_trace(go.Scattermap(
                lat=boundary_lats,
                lon=boundary_lons,
                mode='lines',
                fill='toself',
                fillcolor=color,
                line=dict(width=1, color=color.replace('0.2', '0.5')),
                name=f'Cluster {cluster_id} boundary',
                showlegend=True,
                hoverinfo='skip',
            ))

    # --- Add traces ---

    # Way lines (solid, thicker) — one trace per role
    for role, data in way_lines_by_role.items():
        if data['lats']:
            fig.add_trace(go.Scattermap(
                lat=data['lats'],
                lon=data['lons'],
                mode='lines',
                line=dict(width=3, color=get_role_color(role)),
                name=f'Way: {role}',
                showlegend=True,
                opacity=0.5,
                hoverinfo='skip',
            ))

    # Relation lines (thinner + small dots) — one trace per role
    for role, data in rel_lines_by_role.items():
        if data['lats']:
            fig.add_trace(go.Scattermap(
                lat=data['lats'],
                lon=data['lons'],
                mode='lines+markers',
                line=dict(width=1.5, color=get_role_color(role)),
                marker=dict(size=3, color=get_role_color(role)),
                name=f'Rel: {role}',
                showlegend=True,
                opacity=0.5,
                hoverinfo='skip',
            ))

    # Nodes — one trace per role
    for role, data in nodes_by_role.items():
        if data['lats']:
            fig.add_trace(go.Scattermap(
                lat=data['lats'],
                lon=data['lons'],
                mode='markers',
                marker=dict(size=7, color=get_role_color(role)),
                text=data['texts'],
                name=f'Node: {role}',
                showlegend=True,
                opacity=0.5,
            ))

    # GTFS Stops — always yellow, slightly larger
    if gtfs_lats:
        fig.add_trace(go.Scattermap(
            lat=gtfs_lats,
            lon=gtfs_lons,
            mode='markers',
            marker=dict(size=9, color='rgba(255, 215, 0, 0.6)'),
            text=gtfs_texts,
            name='GTFS Stops',
            opacity=0.6,
        ))

    if all_lats and all_lons:
        center_lat = sum(all_lats) / len(all_lats)
        center_lon = sum(all_lons) / len(all_lons)
    else:
        center_lat = 52.52
        center_lon = 13.40

    fig.update_layout(
        map=dict(
            style="open-street-map",
            center=dict(lat=center_lat, lon=center_lon),
            zoom=10
        ),
        margin={"r":0,"t":0,"l":0,"b":0},
        title="Cluster Visualization"
    )

    # Determine output file path
    if output_path:
        output_file = output_path
    else:
        base_dir = os.path.dirname(os.path.abspath(xml_file))
        output_file = os.path.join(base_dir, "clusters_map.html")

    os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)
    fig.write_html(output_file)
    print(f"Map written to {output_file}")

if __name__ == "__main__":
    
    p = ap.ArgumentParser(description="Plot cluster XML on a map.")
    p.add_argument("xml_file", nargs="?", default="extensive_clusters.xml",
                    help="Path to extensive_clusters.xml[.gz]")
    p.add_argument("--output", "-o", default=None,
                    help="Output HTML file path (default: <xml_dir>/clusters_map.html)")
    args = p.parse_args()
    plot_clusters(args.xml_file, args.output)