CA/backend/routers/risk.py

"""
Router for CBPOA risk assessment endpoints
Reads from GeoJSON files in outputs/daily/ directory
"""
from fastapi import APIRouter, HTTPException, Query
from datetime import datetime, timedelta
from pathlib import Path
from typing import Annotated, List, Literal
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"])


@lru_cache(maxsize=3)
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=3)
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)

    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

        lod_grids = np.column_stack([lat_grid.ravel(), lon_grid.ravel(), risks]).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 for zoom >= 10
    if bounds and zoom >= 10:
        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

    lod_grids = np.column_stack([lat_grid.ravel(), lon_grid.ravel(), risks]).tolist()

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


@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}")

    grids = parse_geojson_file(filepath)

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


@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}")

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

    grids: list[dict[str, str | float]] = []
    for feature in geojson.get("features", []):
        props = feature.get("properties", {})
        coords = feature.get("geometry", {}).get("coordinates", [0, 0])

        risk_value = props.get("risk_1d", 0)
        grids.append({
            "grid_id": str(props.get("node_id", "")),
            "latitude": props.get("lat", coords[1] if len(coords) > 1 else 0),
            "longitude": props.get("lon", coords[0] if len(coords) > 0 else 0),
            "risk_value": risk_value,
            "risk_level": risk_value_to_level(risk_value),
        })

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


@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))
        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()
    )


@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}")

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

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

    return {
        "date": date,
        "total_count": len(grids),
        "columns": ["lat", "lon", "risk_1d", "risk_3d", "risk_7d"],
        "grids": grids,
    }


@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: int = 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}")

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

    target_feature = None
    for feature in geojson.get("features", []):
        props = feature.get("properties", {})
        if str(props.get("node_id", "")) == grid_id:
            target_feature = feature
            break

    if not target_feature 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
        points = []
        features_list = []
        for feature in geojson.get("features", []):
            props = feature.get("properties", {})
            points.append([props.get("lat", 0), props.get("lon", 0)])
            features_list.append(feature)
        if points:
            tree = KDTree(points)
            _, idx = tree.query([center_lat, center_lon])
            target_feature = features_list[idx]

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

    props = target_feature.get("properties", {})
    base_risk = props.get("risk_1d", 0)

    history = []
    for i in range(days):
        history.append({
            "date": (datetime.now() - timedelta(days=i)).strftime("%Y-%m-%d"),
            "risk_value": 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, Query(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_map()

    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(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))
        risk_level = (
            "high" if adjusted >= 0.7 else
            "medium_high" if adjusted >= 0.5 else
            "medium" if adjusted >= 0.3 else
            "medium_low" if adjusted >= 0.2 else
            "low"
        )
        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_level
        ))

    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()
    )