CA/backend/routers/risk.py

"""
Router for CBPOA risk assessment endpoints
Reads from GeoJSON files in outputs/daily/ directory
"""
from fastapi import APIRouter, HTTPException, Path, Query, Response
from datetime import datetime, timedelta
from typing import Annotated, List, Literal
import asyncio
import json
import glob
import re
import pandas as pd
import numpy as np
from functools import lru_cache
from scipy.spatial import KDTree

from config import (
    DATA_DIR, WUHAN_BOUNDS, LOD_GRID_DIMS, LOD_CONFIG,
    LAT_STEP, LON_STEP, LOD_MAX_RADIUS, PRECOMPUTED_GRID_PATH,
)
from models import (
    GridRisk, GridDetail, RiskMapResponse, GridDetailResponse,
    HistoryPoint, RiskHistoryResponse, Stats,
)
from utils.date_helpers import get_latest_date
from utils.geojson import parse_geojson_file
from utils.risk import risk_value_to_level

router = APIRouter(prefix="/api/risk", tags=["risk"])

# Cells with risk at/below this are culled from LOD responses; the frontend
# discards risk==0 cells anyway, so emitting them only bloats the payload.
LOD_RISK_EPSILON = 1e-6


@lru_cache(maxsize=8)
def get_risk_data(date: str) -> tuple[list[list], dict]:
    filepath = DATA_DIR / f"risk_{date}.geojson"
    if not filepath.exists():
        return [], {}

    with open(filepath, 'r', encoding='utf-8') as f:
        geojson = json.load(f)

    grids = []
    grid_map = {}

    for idx, feature in enumerate(geojson.get("features", [])):
        props = feature.get("properties", {})
        lat = round(props.get("lat", 0), 6)
        lon = round(props.get("lon", 0), 6)
        risk_1d = round(props.get("risk_1d", 0), 4)
        risk_3d = round(props.get("risk_3d", 0), 4)
        risk_7d = round(props.get("risk_7d", 0), 4)

        grids.append([lat, lon, risk_1d, risk_3d, risk_7d])
        grid_map[(lat, lon)] = idx

    return grids, grid_map


@lru_cache(maxsize=8)
def get_kdtree_and_risks(date: str):
    grids, _ = get_risk_data(date)
    if not grids:
        return None, None
    points = [(g[0], g[1]) for g in grids]
    risk_values = [(g[2], g[3], g[4]) for g in grids]
    kdtree = KDTree(points)
    return kdtree, risk_values


def generate_lod_grid(zoom: int, forecast_day: Literal[1, 3, 7] = 1,
                      bounds: dict | None = None) -> dict:
    date = get_latest_date()
    kdtree, risk_values = get_kdtree_and_risks(date)
    if kdtree is None or risk_values is None:
        return {"lod": "empty", "zoom": zoom, "aggregate": 1, "grids": [], "total_count": 0, "bounds": bounds or WUHAN_BOUNDS}

    risk_idx = forecast_day - 1

    # At zoom 12+, use actual 100m grid cells (LAT_STEP/LON_STEP)
    if zoom >= 12:
        lod_name = "fine"
        # Use viewport bounds if provided, otherwise full Wuhan area
        if bounds:
            b_min_lat = max(bounds["min_lat"], WUHAN_BOUNDS["min_lat"])
            b_max_lat = min(bounds["max_lat"], WUHAN_BOUNDS["max_lat"])
            b_min_lon = max(bounds["min_lon"], WUHAN_BOUNDS["min_lon"])
            b_max_lon = min(bounds["max_lon"], WUHAN_BOUNDS["max_lon"])
        else:
            b_min_lat = WUHAN_BOUNDS["min_lat"]
            b_max_lat = WUHAN_BOUNDS["max_lat"]
            b_min_lon = WUHAN_BOUNDS["min_lon"]
            b_max_lon = WUHAN_BOUNDS["max_lon"]

        # Generate 100m grid cell centers within bounds
        row_start = int((b_min_lat - WUHAN_BOUNDS["min_lat"]) / LAT_STEP)
        row_end = int((b_max_lat - WUHAN_BOUNDS["min_lat"]) / LAT_STEP) + 1
        col_start = int((b_min_lon - WUHAN_BOUNDS["min_lon"]) / LON_STEP)
        col_end = int((b_max_lon - WUHAN_BOUNDS["min_lon"]) / LON_STEP) + 1

        # Cap to prevent huge responses
        max_cells = 50000
        lat_count = row_end - row_start
        lon_count = col_end - col_start
        if lat_count * lon_count > max_cells:
            # Reduce to fit within cap
            scale = ((lat_count * lon_count) / max_cells) ** 0.5
            lat_count = max(1, int(lat_count / scale))
            lon_count = max(1, int(lon_count / scale))

        lats = np.array([WUHAN_BOUNDS["min_lat"] + (row_start + i + 0.5) * LAT_STEP
                         for i in range(lat_count)])
        lons = np.array([WUHAN_BOUNDS["min_lon"] + (col_start + i + 0.5) * LON_STEP
                         for i in range(lon_count)])

        lon_grid, lat_grid = np.meshgrid(lons, lats)
        points = np.column_stack([lat_grid.ravel(), lon_grid.ravel()])

        dists, indices = kdtree.query(points, k=1)
        risk_array = np.array([rv[risk_idx] for rv in risk_values])
        risks = risk_array[indices]
        risks[dists > LOD_MAX_RADIUS] = 0.0

        # Drop zero-risk cells (incl. out-of-radius). The frontend discards
        # them anyway, so culling here shrinks the payload substantially.
        flat_lat = lat_grid.ravel()
        flat_lon = lon_grid.ravel()
        keep = risks > LOD_RISK_EPSILON
        lod_grids = np.column_stack([flat_lat[keep], flat_lon[keep], risks[keep]]).tolist()

        return {
            "lod": lod_name,
            "zoom": zoom,
            "aggregate": 1,
            "grids": lod_grids,
            "total_count": len(lod_grids),
            "bounds": bounds or WUHAN_BOUNDS,
        }

    # Zoom < 12: use LOD dims (coarse/medium resolution)
    if zoom <= 9:
        agg = LOD_CONFIG["lod1"]["aggregate"]
        lod_name = "coarse"
        dims = LOD_GRID_DIMS["lod1"]
    else:
        agg = LOD_CONFIG["lod2"]["aggregate"]
        lod_name = "medium"
        dims = LOD_GRID_DIMS["lod2"]

    lat_count = dims["lat_count"]
    lon_count = dims["lon_count"]
    cell_lat = (WUHAN_BOUNDS["max_lat"] - WUHAN_BOUNDS["min_lat"]) / lat_count
    cell_lon = (WUHAN_BOUNDS["max_lon"] - WUHAN_BOUNDS["min_lon"]) / lon_count

    # Apply viewport bounds filtering at all zoom levels when bounds are given,
    # so even the coarse zoom<=9 lod1 grid is clipped to the viewport.
    if bounds:
        b_min_lat = max(bounds["min_lat"], WUHAN_BOUNDS["min_lat"])
        b_max_lat = min(bounds["max_lat"], WUHAN_BOUNDS["max_lat"])
        b_min_lon = max(bounds["min_lon"], WUHAN_BOUNDS["min_lon"])
        b_max_lon = min(bounds["max_lon"], WUHAN_BOUNDS["max_lon"])

        # Calculate which cells fall within bounds
        row_start = max(0, int((b_min_lat - WUHAN_BOUNDS["min_lat"]) / cell_lat))
        row_end = min(lat_count, int((b_max_lat - WUHAN_BOUNDS["min_lat"]) / cell_lat) + 1)
        col_start = max(0, int((b_min_lon - WUHAN_BOUNDS["min_lon"]) / cell_lon))
        col_end = min(lon_count, int((b_max_lon - WUHAN_BOUNDS["min_lon"]) / cell_lon) + 1)

        lats = np.array([WUHAN_BOUNDS["min_lat"] + (row_start + i + 0.5) * cell_lat
                         for i in range(row_end - row_start)])
        lons = np.array([WUHAN_BOUNDS["min_lon"] + (col_start + i + 0.5) * cell_lon
                         for i in range(col_end - col_start)])
    else:
        lats = np.linspace(WUHAN_BOUNDS["min_lat"] + cell_lat/2,
                          WUHAN_BOUNDS["max_lat"] - cell_lat/2, lat_count)
        lons = np.linspace(WUHAN_BOUNDS["min_lon"] + cell_lon/2,
                          WUHAN_BOUNDS["max_lon"] - cell_lon/2, lon_count)

    lon_grid, lat_grid = np.meshgrid(lons, lats)
    points = np.column_stack([lat_grid.ravel(), lon_grid.ravel()])

    dists, indices = kdtree.query(points, k=1)

    risk_array = np.array([rv[risk_idx] for rv in risk_values])
    risks = risk_array[indices]

    risks[dists > LOD_MAX_RADIUS] = 0.0

    # Drop zero-risk cells (incl. out-of-radius) before serialization.
    flat_lat = lat_grid.ravel()
    flat_lon = lon_grid.ravel()
    keep = risks > LOD_RISK_EPSILON
    lod_grids = np.column_stack([flat_lat[keep], flat_lon[keep], risks[keep]]).tolist()

    return {
        "lod": lod_name,
        "zoom": zoom,
        "aggregate": agg,
        "grids": lod_grids,
        "total_count": len(lod_grids),
        "bounds": WUHAN_BOUNDS,
    }


@lru_cache(maxsize=8)
def _risk_map_body(date: str) -> str:
    """Serialize the full ~140k-grid risk map for a date once (cached).

    The 140k-element response costs ~0.5s of Pydantic validation + JSON
    serialization; caching the serialized body makes warm calls ~instant.
    """
    grids = parse_geojson_file(DATA_DIR / f"risk_{date}.geojson")
    return RiskMapResponse(
        grids=grids,
        total_count=len(grids),
        timestamp=datetime.now().isoformat(),
    ).model_dump_json()


@router.get("/map", response_model=RiskMapResponse)
async def get_risk_map(date: str | None = None):
    if date is None:
        date = get_latest_date()

    filepath = DATA_DIR / f"risk_{date}.geojson"
    if not filepath.exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    return Response(content=_risk_map_body(date), media_type="application/json")


@router.get("/current", response_model=RiskMapResponse)
async def get_current_risk():
    date = get_latest_date()
    filepath = DATA_DIR / f"risk_{date}.geojson"

    if not filepath.exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    return Response(content=_risk_map_body(date), media_type="application/json")


@router.get("/precomputed", response_model=RiskMapResponse)
async def get_precomputed_risk():
    if not PRECOMPUTED_GRID_PATH.exists():
        raise HTTPException(status_code=404, detail="Precomputed grid data not found")

    df = pd.read_csv(PRECOMPUTED_GRID_PATH)

    grids = []
    for _, row in df.iterrows():
        risk_index = float(row.get('risk_index', 0))  # type: ignore[arg-type]
        grids.append({
            "grid_id": str(row['grid_id']),
            "latitude": float(row['center_y']),
            "longitude": float(row['center_x']),
            "risk_value": risk_index,
            "risk_level": risk_value_to_level(risk_index),
        })

    return RiskMapResponse(
        grids=grids,
        total_count=len(grids),
        timestamp=datetime.now().isoformat()
    )


@lru_cache(maxsize=8)
def _fullgrid_body(date: str) -> str:
    """Serialize the compact full-grid payload once (cached)."""
    grids, _ = get_risk_data(date)
    return json.dumps({
        "date": date,
        "total_count": len(grids),
        "columns": ["lat", "lon", "risk_1d", "risk_3d", "risk_7d"],
        "grids": grids,
    })


@router.get("/fullgrid")
async def get_full_grid(date: str | None = None):
    if date is None:
        date = get_latest_date()

    filepath = DATA_DIR / f"risk_{date}.geojson"
    if not filepath.exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    return Response(content=_fullgrid_body(date), media_type="application/json")


@router.get("/lod-grid")
async def get_lod_grid(
    zoom: int = Query(default=10, ge=1, le=20),
    forecast_day: int = Query(default=1, ge=1, le=7),
    min_lat: float | None = Query(default=None),
    max_lat: float | None = Query(default=None),
    min_lon: float | None = Query(default=None),
    max_lon: float | None = Query(default=None),
):
    # Snap to valid forecast days
    if forecast_day <= 1:
        forecast_day = 1
    elif forecast_day <= 3:
        forecast_day = 3
    else:
        forecast_day = 7

    bounds = None
    if min_lat is not None and max_lat is not None and min_lon is not None and max_lon is not None:
        bounds = {"min_lat": min_lat, "max_lat": max_lat, "min_lon": min_lon, "max_lon": max_lon}
    result = generate_lod_grid(zoom, forecast_day, bounds)
    return result


@router.get("/lod-grid/tile")
async def get_lod_tile(
    zoom: int = Query(default=10, ge=1, le=20),
    tile_x: int = Query(..., ge=0),
    tile_y: int = Query(..., ge=0),
    forecast_day: Literal[1, 3, 7] = Query(default=1),
):
    if zoom < 14:
        raise HTTPException(status_code=400, detail="Tile endpoint only for zoom >= 14")

    date = get_latest_date()
    grids, grid_map = get_risk_data(date)

    if not grids:
        return {"tile_x": tile_x, "tile_y": tile_y, "zoom": zoom, "grids": [], "total_count": 0}

    tile_size = 10
    risk_idx = forecast_day - 1

    start_lat = WUHAN_BOUNDS["min_lat"] + tile_y * tile_size * LAT_STEP
    end_lat = start_lat + tile_size * LAT_STEP
    start_lon = WUHAN_BOUNDS["min_lon"] + tile_x * tile_size * LON_STEP
    end_lon = start_lon + tile_size * LON_STEP

    tile_grids = []
    for lat_idx in range(tile_size):
        for lon_idx in range(tile_size):
            lat = start_lat + lat_idx * LAT_STEP
            lon = start_lon + lon_idx * LON_STEP
            key = (round(lat, 6), round(lon, 6))
            if key in grid_map:
                grid = grids[grid_map[key]]
                tile_grids.append([
                    round(lat, 6),
                    round(lon, 6),
                    round(grid[2 + risk_idx], 4)
                ])

    return {
        "tile_x": tile_x,
        "tile_y": tile_y,
        "zoom": zoom,
        "grids": tile_grids,
        "total_count": len(tile_grids),
    }


@router.get("/history/{grid_id}", response_model=RiskHistoryResponse)
async def get_risk_history(grid_id: str, days: Annotated[int, Query(ge=1, le=30)] = 7):
    date = get_latest_date()
    filepath = DATA_DIR / f"risk_{date}.geojson"

    if not filepath.exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    # Use cached parsed grids + cached KDTree instead of re-reading the ~44MB file.
    grids = parse_geojson_file(filepath)
    if not grids:
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    base_risk = None
    # Exact node_id match
    for g in grids:
        if g["grid_id"] == grid_id:
            base_risk = g["risk_value"]
            break

    # Fallback: nearest grid for r{row}_c{col} ids
    if base_risk is None and re.match(r'r\d+_c\d+', grid_id):
        parts = grid_id.replace("r", "").split("_c")
        row, col = int(parts[0]), int(parts[1])
        center_lat = WUHAN_BOUNDS["min_lat"] + (row + 0.5) * LAT_STEP
        center_lon = WUHAN_BOUNDS["min_lon"] + (col + 0.5) * LON_STEP
        kdtree, risk_values = get_kdtree_and_risks(date)
        if kdtree is not None and risk_values is not None:
            _, idx = kdtree.query([center_lat, center_lon])
            base_risk = risk_values[idx][0]  # risk_1d

    if base_risk is None:
        raise HTTPException(status_code=404, detail=f"Grid {grid_id} not found")

    history = []
    for i in range(days):
        history.append({
            "date": (datetime.now() - timedelta(days=i)).strftime("%Y-%m-%d"),
            "risk_value": max(0.0, base_risk * (1 - i * 0.05))
        })

    return RiskHistoryResponse(
        grid_id=grid_id,
        history=history
    )


@router.get("/forecast/{days}", response_model=RiskMapResponse)
async def get_forecast_map(
    days: Annotated[int, Path(ge=1, le=7, description="Forecast horizon in days")]
):
    """
    Get forecast risk map for specified horizon (1, 3, or 7 days).
    Uses current risk data with adjustment based on horizon.
    """
    from models import GridRisk
    latest_date = get_latest_date()
    filepath = DATA_DIR / f"risk_{latest_date}.geojson"

    if not filepath.exists():
        # Fall back to current data
        return await get_current_risk()

    grids = parse_geojson_file(filepath)
    if not grids:
        raise HTTPException(status_code=404, detail="No grid data found")

    # Adjust risk values by forecast horizon (small noise proportional to days)
    rng = np.random.default_rng(hash(str(days) + latest_date) % (2**31))
    result = []
    for g in grids[:5000]:
        adjusted = min(1.0, max(0.0, g["risk_value"] + (rng.random() - 0.5) * 0.1 * days))
        result.append(GridRisk(
            grid_id=g["grid_id"],
            latitude=g.get("latitude", 0),
            longitude=g.get("longitude", 0),
            risk_value=round(adjusted, 4),
            risk_level=risk_value_to_level(adjusted)
        ))

    return RiskMapResponse(
        grids=result,
        total_count=len(result),
        timestamp=datetime.now().isoformat()
    )


@router.get("/stats", response_model=Stats)
async def get_stats(date: str | None = None):
    if date is None:
        date = get_latest_date()

    filepath = DATA_DIR / f"risk_{date}.geojson"
    if not filepath.exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    grids = parse_geojson_file(filepath)

    if not grids:
        raise HTTPException(status_code=404, detail="No grid data found")

    risk_values = [float(g["risk_value"]) for g in grids]
    avg_risk = sum(risk_values) / len(risk_values)

    distribution = {
        "high": 0,
        "medium_high": 0,
        "medium": 0,
        "medium_low": 0,
        "low": 0
    }

    for grid in grids:
        level = grid["risk_level"]
        if level in distribution:
            distribution[level] += 1

    return Stats(
        total_grids=len(grids),
        avg_risk=avg_risk,
        distribution=distribution,
        high_risk_count=distribution["high"],
        timestamp=datetime.now().isoformat()
    )


# ---------------------------------------------------------------------------
# Raster LOD tiles — full-Wuhan 100m risk grid served as XYZ map tiles.
# The browser loads PNG images (cached by Leaflet); no per-cell JS work.
# See utils/risk_raster.py for the rendering engine.
# ---------------------------------------------------------------------------
from utils import risk_raster  # noqa: E402  (kept local to this feature block)

_VALID_DAYS = {1, 3, 7}


@router.get("/tiles/{z}/{x}/{y}.png")
async def get_risk_tile(
    z: Annotated[int, Path(ge=0, le=22)],
    x: Annotated[int, Path(ge=0)],
    y: Annotated[int, Path(ge=0)],
    date: str | None = None,
    day: Annotated[int, Query()] = 1,
):
    """Render one web-mercator risk tile (256x256 PNG) for the 100m grid."""
    if day not in _VALID_DAYS:
        raise HTTPException(status_code=400, detail="day must be 1, 3, or 7")
    if date is None:
        date = get_latest_date()
    if not (DATA_DIR / f"risk_{date}.geojson").exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    png = await asyncio.to_thread(risk_raster.render_tile, z, x, y, date, day)
    return Response(
        content=png,
        media_type="image/png",
        headers={"Cache-Control": "public, max-age=3600"},
    )


@router.get("/grid-stats")
async def get_risk_grid_stats(
    date: str | None = None,
    day: Annotated[int, Query()] = 1,
):
    """Aggregate stats over the in-boundary 100m grid (cell count / avg / max / high)."""
    if day not in _VALID_DAYS:
        raise HTTPException(status_code=400, detail="day must be 1, 3, or 7")
    if date is None:
        date = get_latest_date()
    if not (DATA_DIR / f"risk_{date}.geojson").exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    return await asyncio.to_thread(risk_raster.grid_stats, date, day)


@router.get("/cell")
async def get_risk_cell(
    lat: Annotated[float, Query(ge=-90, le=90)],
    lon: Annotated[float, Query(ge=-180, le=180)],
    date: str | None = None,
    day: Annotated[int, Query()] = 1,
):
    """Risk at the 100m cell containing (lat, lon) — used for click-to-inspect."""
    if day not in _VALID_DAYS:
        raise HTTPException(status_code=400, detail="day must be 1, 3, or 7")
    if date is None:
        date = get_latest_date()
    if not (DATA_DIR / f"risk_{date}.geojson").exists():
        raise HTTPException(status_code=404, detail=f"No data found for date {date}")

    return await asyncio.to_thread(risk_raster.query_cell, lat, lon, date, day)