src/dashboard/components/pydeck_layers.py

"""
Pydeck layer factories — shared across Subsystems K (3D terrain & surge),
L (cyclone track & probability cone), and M (evacuation flow network).

Each function returns a ``pydeck.Layer`` configured with Aegis Command
colors (defined in ``src.config``). The module also exposes
``get_map_style()`` which picks between Mapbox dark-v11 (if a token is
present in ``.streamlit/secrets.toml``) and the CartoDB dark-matter fallback.
"""
from __future__ import annotations

import pydeck as pdk
import pandas as pd
import geopandas as gpd

from src.config import (
    DECK_SEVERITY_COLORS,
    DECK_COLOR_SHELTER,
    DECK_COLOR_SCHOOL,
    DECK_COLOR_SCHOOL_IN_ZONE,
    DECK_COLOR_TOWER_OK,
    DECK_COLOR_TOWER_DOWN,
    DECK_COLOR_ROAD_OK,
    DECK_COLOR_ROAD_CUT,
    get_mapbox_token,
)


# ─── Basemap Style ────────────────────────────────────────────────────────────
def get_map_style() -> str:
    """Return a dark basemap style URI.

    Prefers Mapbox ``dark-v11`` when a token is configured; otherwise falls
    back to the free CartoDB *dark-matter* style which requires no auth.
    """
    token = get_mapbox_token()
    if token:
        return "mapbox://styles/mapbox/dark-v11"
    return "https://basemaps.cartocdn.com/gl/dark-matter-gl-style/style.json"


# ─── 3D Terrain Base ──────────────────────────────────────────────────────────
def make_terrain_layer(bounds: list[float]) -> pdk.Layer:
    """3D terrain using the free Mapzen/AWS Terrarium elevation tiles.

    Parameters
    ----------
    bounds : [west, south, east, north]
        Decimal-degree bounding box the terrain is rendered within.
    """
    return pdk.Layer(
        "TerrainLayer",
        elevation_decoder={
            "rScaler": 256,
            "gScaler": 1,
            "bScaler": 1 / 256,
            "offset": -32768,
        },
        elevation_data="https://s3.amazonaws.com/elevation-tiles-prod/terrarium/{z}/{x}/{y}.png",
        texture=(
            "https://server.arcgisonline.com/ArcGIS/rest/services/"
            "World_Imagery/MapServer/tile/{z}/{y}/{x}"
        ),
        bounds=bounds,
        material={
            "ambient": 0.2,
            "diffuse": 0.6,
            "shininess": 32,
            "specularColor": [60, 64, 70],
        },
    )


# ─── Severity Choropleth with Extrusion ──────────────────────────────────────
def make_impact_extrusion(
    impact_geojson: dict,
    elevation_key: str = "surge_depth_m",
    elevation_scale: float = 800.0,
) -> pdk.Layer:
    """Extruded GeoJsonLayer — surge depth drives polygon height.

    The caller **must** have annotated each feature with numeric
    ``severity_rgb_r/g/b`` properties (see ``annotate_severity_rgb``
    helper in the consuming page). Pydeck resolves the colour
    expression against those numeric properties.
    """
    return pdk.Layer(
        "GeoJsonLayer",
        data=impact_geojson,
        pickable=True,
        stroked=True,
        filled=True,
        extruded=True,
        wireframe=True,
        get_fill_color=(
            "[properties.severity_rgb_r, properties.severity_rgb_g, "
            "properties.severity_rgb_b, 140]"
        ),
        get_line_color=[148, 163, 184, 200],
        get_elevation=f"properties.{elevation_key} * {elevation_scale}",
        line_width_min_pixels=1,
    )


# ─── Shelter Columns ──────────────────────────────────────────────────────────
def make_shelter_columns(shelter_df: pd.DataFrame) -> pdk.Layer:
    """ColumnLayer — height proportional to shelter capacity.

    ``shelter_df`` must have columns: ``lon``, ``lat``, ``capacity``, ``name``.
    """
    return pdk.Layer(
        "ColumnLayer",
        data=shelter_df,
        get_position=["lon", "lat"],
        get_elevation="capacity",
        elevation_scale=2.5,
        radius=150,
        get_fill_color=DECK_COLOR_SHELTER,
        pickable=True,
        auto_highlight=True,
    )


# ─── Schools — Scatterplot (red inside hazard zones) ─────────────────────────
def make_school_scatter(school_df: pd.DataFrame) -> pdk.Layer:
    """ScatterplotLayer — schools coloured cyan (safe) or red (in hazard zone).

    ``school_df`` must have columns: ``lon``, ``lat``, ``name``,
    ``in_hazard_zone`` (bool), ``protocol_status`` (str, may be empty).
    """
    school_df = school_df.copy()
    school_df["color"] = school_df["in_hazard_zone"].apply(
        lambda x: DECK_COLOR_SCHOOL_IN_ZONE if x else DECK_COLOR_SCHOOL
    )
    return pdk.Layer(
        "ScatterplotLayer",
        data=school_df,
        get_position=["lon", "lat"],
        get_fill_color="color",
        get_radius=200,
        radius_min_pixels=3,
        radius_max_pixels=10,
        pickable=True,
        auto_highlight=True,
    )


# ─── Cyclone Track Arcs (Subsystem L) ────────────────────────────────────────
def make_track_arcs(
    track_df: pd.DataFrame,
    highlight_idx: int | None = None,
) -> list[pdk.Layer]:
    """Render a cyclone track as coloured arc segments + a pulsing current point.

    Parameters
    ----------
    track_df : DataFrame with columns ``timestamp``, ``lat``, ``lon``, ``wind_knots``.
    highlight_idx : int or None — index of the "current" track point.
        Defaults to the last row.

    Returns
    -------
    list[pdk.Layer]
        Two layers: (1) an ArcLayer for segments, (2) a ScatterplotLayer pulse
        for the current position. Caller spreads these into the deck's layers.
    """
    segs: list[dict] = []
    for i in range(len(track_df) - 1):
        w = track_df.iloc[i].get("wind_knots") or 0
        color = _wind_color(float(w))
        segs.append({
            "from_lon": float(track_df.iloc[i]["lon"]),
            "from_lat": float(track_df.iloc[i]["lat"]),
            "to_lon":   float(track_df.iloc[i + 1]["lon"]),
            "to_lat":   float(track_df.iloc[i + 1]["lat"]),
            "color_r": color[0], "color_g": color[1], "color_b": color[2],
            "width": 2 + (float(w) / 20.0),
        })

    arcs = pdk.Layer(
        "ArcLayer",
        data=segs,
        get_source_position=["from_lon", "from_lat"],
        get_target_position=["to_lon", "to_lat"],
        get_source_color="[color_r, color_g, color_b, 200]",
        get_target_color="[color_r, color_g, color_b, 200]",
        get_width="width",
        pickable=False,
    )

    current_pos = (
        track_df.iloc[highlight_idx] if highlight_idx is not None
        else track_df.iloc[-1]
    )
    pulse = pdk.Layer(
        "ScatterplotLayer",
        data=[{
            "lon": float(current_pos["lon"]),
            "lat": float(current_pos["lat"]),
        }],
        get_position=["lon", "lat"],
        get_fill_color=[255, 49, 49, 220],
        get_radius=8000,
        radius_min_pixels=8,
        radius_max_pixels=30,
        pickable=False,
    )
    return [arcs, pulse]


def _wind_color(wind_knots: float) -> list[int]:
    """Saffir-Simpson-style colour ramp: green → yellow → orange → red."""
    if wind_knots < 34:
        return [0, 255, 157]     # Tropical Depression
    if wind_knots < 64:
        return [255, 184, 0]     # Tropical Storm
    if wind_knots < 83:
        return [255, 107, 53]    # Category 1
    if wind_knots < 96:
        return [255,  70, 30]    # Category 2
    if wind_knots < 113:
        return [255,  49, 49]    # Category 3
    return [200, 0, 60]          # Category 4+


# ─── Probability Cone Heatmap (Subsystem L) ──────────────────────────────────
def make_probability_heatmap(
    ensemble_points: pd.DataFrame,
    weight_key: str = "prob",
) -> pdk.Layer:
    """HeatmapLayer rendering the ensemble track probability cone.

    ``ensemble_points`` must have columns: ``lon``, ``lat``, and a numeric
    weight column (default ``prob``). Suitable input is the flattened
    500-member ensemble produced by Part II Subsystem C.
    """
    return pdk.Layer(
        "HeatmapLayer",
        data=ensemble_points,
        get_position=["lon", "lat"],
        get_weight=weight_key,
        radiusPixels=60,
        intensity=1.2,
        threshold=0.08,
        color_range=[
            [  0, 240, 255,   0],
            [  0, 240, 255, 100],
            [255, 184,   0, 160],
            [255, 107,  53, 200],
            [255,  49,  49, 240],
        ],
    )


# ─── Road Network (Subsystem M) ──────────────────────────────────────────────
def make_road_paths(road_gdf: gpd.GeoDataFrame) -> pdk.Layer:
    """PathLayer — roads coloured by passability.

    ``road_gdf`` must contain LineString geometries and a ``passable``
    boolean column. An optional ``highway`` column controls line width
    (primary/trunk = thicker).
    """
    road_records: list[dict] = []
    for _, row in road_gdf.iterrows():
        geom = row.geometry
        if geom is None or geom.is_empty:
            continue
        # Support MultiLineString by flattening to individual paths.
        if geom.geom_type == "MultiLineString":
            lines = list(geom.geoms)
        else:
            lines = [geom]
        for line in lines:
            coords = list(line.coords)
            if len(coords) < 2:
                continue
            road_records.append({
                "path": [[c[0], c[1]] for c in coords],
                "color": (
                    DECK_COLOR_ROAD_OK if row.get("passable", True)
                    else DECK_COLOR_ROAD_CUT
                ),
                "width": 2 if row.get("highway", "") in ("primary", "trunk") else 1,
            })

    return pdk.Layer(
        "PathLayer",
        data=road_records,
        get_path="path",
        get_color="color",
        get_width="width",
        width_min_pixels=1,
        width_max_pixels=4,
        pickable=False,
    )


# ─── Evacuation Flow Lines (Subsystem M) ─────────────────────────────────────
def make_evacuation_flows(assignments_df: pd.DataFrame) -> pdk.Layer:
    """LineLayer — population-weighted flows from origin centroids to shelters.

    ``assignments_df`` columns:
        ``origin_lon``, ``origin_lat``, ``dest_lon``, ``dest_lat``,
        ``population``, ``status`` — one of
        ``accessible`` / ``long`` / ``severed``.
    """
    status_colors = {
        "accessible": [  0, 255, 157, 200],
        "long":       [255, 184,   0, 200],
        "severed":    [255,  49,  49, 220],
    }

    assignments_df = assignments_df.copy()
    assignments_df["color"] = assignments_df["status"].apply(
        lambda s: status_colors.get(s, [148, 163, 184, 180])
    )
    pop_max = max(float(assignments_df["population"].max() or 1), 1.0)
    assignments_df["width"] = (
        assignments_df["population"].astype(float) / pop_max * 8 + 1
    )

    return pdk.Layer(
        "LineLayer",
        data=assignments_df,
        get_source_position=["origin_lon", "origin_lat"],
        get_target_position=["dest_lon", "dest_lat"],
        get_color="color",
        get_width="width",
        pickable=True,
    )


# ─── Telecom Towers (Warning Dispatch page) ──────────────────────────────────
def make_tower_layer(tower_df: pd.DataFrame) -> pdk.Layer:
    """ScatterplotLayer — telecom towers (green=operational, red=down).

    ``tower_df``: ``lon``, ``lat``, ``operational`` (bool), ``radio_type``.
    """
    tower_df = tower_df.copy()
    tower_df["color"] = tower_df["operational"].apply(
        lambda x: DECK_COLOR_TOWER_OK if x else DECK_COLOR_TOWER_DOWN
    )
    return pdk.Layer(
        "ScatterplotLayer",
        data=tower_df,
        get_position=["lon", "lat"],
        get_fill_color="color",
        get_radius=180,
        radius_min_pixels=2,
        radius_max_pixels=6,
        pickable=True,
    )


# ─── Severity RGB Annotation Helper ──────────────────────────────────────────
def annotate_severity_rgb(geojson: dict) -> dict:
    """Inject numeric ``severity_rgb_{r,g,b}`` properties into each feature.

    Pydeck's JSON config evaluates ``get_fill_color`` expressions against
    numeric feature properties more reliably than string severity labels.
    This helper mutates and returns the same GeoJSON object.
    """
    for feat in geojson.get("features", []):
        sev = feat.get("properties", {}).get("severity", "UNAFFECTED")
        r, g, b, _ = DECK_SEVERITY_COLORS.get(sev, DECK_SEVERITY_COLORS["UNAFFECTED"])
        feat["properties"]["severity_rgb_r"] = r
        feat["properties"]["severity_rgb_g"] = g
        feat["properties"]["severity_rgb_b"] = b
        feat["properties"].setdefault("surge_depth_m", 0.0)
    return geojson