CA/backend/routers/grid.py

import asyncio
from fastapi import APIRouter, HTTPException, Query, Response
from datetime import datetime, timedelta
from functools import lru_cache
from pathlib import Path
from typing import Optional
import logging
import sys
import math
import pandas as pd

PROJECT_ROOT = Path(__file__).parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from models import (
    DistrictAggregation,
    HistoricalAggregationRequest,
    HistoricalAggregationResponse,
    GridGeoJSONResponse,
    GridPrediction,
    MultiDayPredictionRequest,
    MultiDayPredictionResponse,
)
from data.case_loader import load_cases_by_district_daily

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

logger = logging.getLogger("cbpoa.grid")

_parquet_cache: dict[str, pd.DataFrame] = {}

def _load_parquet(path: Path) -> pd.DataFrame:
    key = str(path)
    if key not in _parquet_cache:
        _parquet_cache[key] = pd.read_parquet(path)
    return _parquet_cache[key]


def _compute_historical_aggregation(
    start: datetime,
    end: datetime,
    aggregation: str,
    district: Optional[str],
) -> HistoricalAggregationResponse:
    """Run the full pandas aggregation pipeline (called in thread pool)."""
    try:
        cases_df = load_cases_by_district_daily()
    except FileNotFoundError:
        return HistoricalAggregationResponse(
            aggregations=[], total_records=0,
            date_range=(start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")),
            timestamp=datetime.now().isoformat(),
        )
    cases_df['date'] = pd.to_datetime(cases_df['date'])

    filtered_cases = cases_df[
        (cases_df['date'] >= start) &
        (cases_df['date'] <= end)
    ]

    if district:
        filtered_cases = filtered_cases[
            filtered_cases['district'].str.contains(district.replace('区', ''), na=False, regex=False)
        ]

    if aggregation == "weekly":
        filtered_cases = filtered_cases.copy()
        filtered_cases['period'] = filtered_cases['date'].dt.to_period('W').astype(str)
        grouped = filtered_cases.groupby(['period', 'district']).agg({
            'total_cases': 'sum',
            'outpatient_count': 'sum',
            'inpatient_count': 'sum',
        }).reset_index()
        grouped['date'] = grouped['period']
    elif aggregation == "monthly":
        filtered_cases = filtered_cases.copy()
        filtered_cases['period'] = filtered_cases['date'].dt.to_period('M').astype(str)
        grouped = filtered_cases.groupby(['period', 'district']).agg({
            'total_cases': 'sum',
            'outpatient_count': 'sum',
            'inpatient_count': 'sum',
        }).reset_index()
        grouped['date'] = grouped['period']
    else:
        grouped = filtered_cases.copy()
        grouped['date'] = grouped['date'].dt.strftime('%Y-%m-%d')

    try:
        weather_df = _load_parquet(PROJECT_ROOT / "processed" / "weather" / "station_daily_2022.parquet")
    except FileNotFoundError:
        weather_df = pd.DataFrame({'date': pd.Series(dtype='str'), 'AQI': pd.Series(dtype='float64'), 'PM25': pd.Series(dtype='float64'), 'PM10': pd.Series(dtype='float64')})
    weather_df = weather_df.copy()
    weather_df['date'] = pd.to_datetime(weather_df['date']).dt.strftime('%Y-%m-%d')

    # Weather data doesn't have district - aggregate by date only
    weather_agg = weather_df.groupby(['date']).agg({
        'AQI': 'mean',
        'PM25': 'mean',
        'PM10': 'mean',
    }).reset_index()

    # Merge by date only
    merged = grouped.merge(weather_agg, on=['date'], how='left')

    aggregations = []
    for _, row in merged.iterrows():
        aggregations.append(DistrictAggregation(
            district=str(row['district']),
            date=str(row['date']),
            total_cases=int(row['total_cases']),
            outpatient_count=int(row['outpatient_count']),
            inpatient_count=int(row['inpatient_count']),
            avg_AQI=float(row['AQI']) if bool(pd.notna(row['AQI'])) else 0.0,
            avg_PM25=float(row['PM25']) if bool(pd.notna(row['PM25'])) else 0.0,
            avg_PM10=float(row['PM10']) if bool(pd.notna(row['PM10'])) else 0.0,
        ))

    return HistoricalAggregationResponse(
        aggregations=aggregations,
        total_records=len(aggregations),
        date_range=(start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")),
        timestamp=datetime.now().isoformat(),
    )


@router.get("/history/aggregated", response_model=HistoricalAggregationResponse)
async def get_historical_aggregated(
    start_date: str = Query(..., description="Start date (YYYY-MM-DD)"),
    end_date: str = Query(..., description="End date (YYYY-MM-DD)"),
    aggregation: str = Query("daily", description="Aggregation level: daily, weekly, monthly"),
    district: Optional[str] = Query(None, description="Filter by district name"),
):
    """
    Historical data aggregation API.

    Returns aggregated case and weather data by district and date.
    Pandas processing runs in a thread pool to avoid blocking the async event loop.
    """
    try:
        start = datetime.strptime(start_date, "%Y-%m-%d")
        end = datetime.strptime(end_date, "%Y-%m-%d")
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid date format. Use YYYY-MM-DD")

    if (end - start).days > 365:
        raise HTTPException(status_code=400, detail="Date range exceeds 365 days")

    # Offload all pandas I/O and processing to a thread pool
    # to prevent blocking the async event loop
    return await asyncio.to_thread(
        _compute_historical_aggregation, start, end, aggregation, district
    )


@lru_cache(maxsize=1)
def _grid_geojson_base():
    """Date-independent base merge: grid centroid + district + real population.

    Merged once and cached (the source frames are ~1M rows each, so the join
    must not run per request). Raises FileNotFoundError if the core grid files
    are missing (caller handles it).
    """
    import pandas as pd
    grid_df = _load_parquet(PROJECT_ROOT / "processed" / "grid_100m_index.parquet")
    district_map = _load_parquet(PROJECT_ROOT / "processed" / "grid_district_mapping.parquet")
    base = grid_df.merge(district_map, on='grid_id', how='left')
    try:
        pop_df = _load_parquet(PROJECT_ROOT / "processed" / "grid_100m_with_dem_pop.parquet")
        base = base.merge(pop_df[['grid_id', 'population_density']], on='grid_id', how='left')
    except FileNotFoundError:
        base['population_density'] = 0.0
    base['population_density'] = base['population_density'].fillna(0.0)
    return base


def _risk_level_of(v: float) -> str:
    if v >= 0.7:
        return "high"
    if v >= 0.5:
        return "medium"
    if v >= 0.3:
        return "medium_low"
    return "low"


@lru_cache(maxsize=32)
def _grids_geojson_body(date: str, district: Optional[str], risk_level: Optional[str]) -> str:
    """Build + serialize the grid GeoJSON once per (date, district, risk_level).

    Risk is computed vectorised over the full grid (no per-row Python loop) and
    the highest-risk grids are returned as hotspots, so the map shows real
    high→low variation. Cached, so warm calls are near-instant. Raises
    FileNotFoundError if the core grid files are missing.
    """
    merged = _grid_geojson_base()
    if district:
        merged = merged[merged['district_name'].str.contains(district.replace('区', ''), na=False, regex=False)]

    cases_df = load_cases_by_district_daily()
    cases_df['date'] = pd.to_datetime(cases_df['date']).dt.strftime('%Y-%m-%d')
    cases_df = cases_df[cases_df['date'] == date]

    # Normalise district case load to 0..1 across districts for this date.
    max_district_cases = float(cases_df['total_cases'].max()) if len(cases_df) else 0.0
    if max_district_cases <= 0:
        max_district_cases = 1.0

    merged = merged.merge(cases_df[['district', 'total_cases']], left_on='district_name', right_on='district', how='left')
    merged = merged.copy()
    merged['total_cases'] = merged['total_cases'].fillna(0).astype(int)
    merged['center_lon'] = pd.to_numeric(merged['center_lon'], errors='coerce').fillna(0.0)
    merged['center_lat'] = pd.to_numeric(merged['center_lat'], errors='coerce').fillna(0.0)
    merged['population_density'] = merged['population_density'].fillna(0.0).clip(lower=0.0)

    # Drop grids without coordinates.
    merged = merged[(merged['center_lon'] != 0.0) | (merged['center_lat'] != 0.0)]

    # Demo risk model (vectorised): a district's relative case load × each grid's
    # own population exposure. Sparse cells stay low; densely-populated cells in
    # high-case districts rise toward 1.0.
    district_load = (merged['total_cases'] / max_district_cases).clip(upper=1.0)
    pop_factor = (merged['population_density'] / 50.0).clip(upper=1.0)
    merged['risk_value'] = (0.1 + 0.85 * district_load * pop_factor).clip(upper=1.0).round(3)

    # Show the highest-risk grids (hotspots), not arbitrary cells.
    merged = merged.nlargest(10000, 'risk_value')

    features = []
    for rec in merged.to_dict('records'):
        rv = float(rec['risk_value'])
        lvl = _risk_level_of(rv)
        if risk_level and lvl != risk_level:
            continue
        name = rec.get('district_name')
        if not isinstance(name, str):
            name = "未知"
        lon = round(float(rec['center_lon']), 6)
        lat = round(float(rec['center_lat']), 6)
        features.append({
            "type": "Feature",
            "geometry": {"type": "Point", "coordinates": [lon, lat]},
            "properties": {
                "grid_id": str(rec.get('grid_id', '')),
                "latitude": lat,
                "longitude": lon,
                "district": name,
                "total_cases": int(rec.get('total_cases', 0)),
                "population_density": round(float(rec.get('population_density', 0.0)), 2),
                "risk_value": rv,
                "risk_level": lvl,
            }
        })

    return GridGeoJSONResponse(
        type="FeatureCollection",
        features=features,
        timestamp=datetime.now().isoformat(),
    ).model_dump_json()


@router.get("/grids/geojson", response_model=GridGeoJSONResponse)
async def get_grids_geojson(
    date: str = Query(..., description="Date (YYYY-MM-DD)"),
    district: Optional[str] = Query(None, description="Filter by district"),
    risk_level: Optional[str] = Query(None, description="Filter by risk level"),
):
    """Get grid data as GeoJSON for map visualization (cached per query)."""
    try:
        # Offload the parquet merges + vectorised compute to a thread so the
        # cold-cache build doesn't block the event loop.
        body = await asyncio.to_thread(_grids_geojson_body, date, district, risk_level)
    except FileNotFoundError:
        return GridGeoJSONResponse(type="FeatureCollection", features=[], timestamp=datetime.now().isoformat())
    return Response(content=body, media_type="application/json")


@router.post("/predict/multi-day", response_model=MultiDayPredictionResponse)
async def predict_multi_day(request: MultiDayPredictionRequest):
    """
    Multi-day prediction API for grid-level risk assessment.

    Returns risk predictions for each grid cell across multiple days.
    Uses the SpatialTemporalGCN model with on-demand feature generation.
    """
    from scripts.generate_grid_features import GridFeatureGenerator

    try:
        start_date = datetime.strptime(request.date, "%Y-%m-%d")
    except ValueError:
        raise HTTPException(status_code=400, detail="Invalid date format. Use YYYY-MM-DD")

    generator = GridFeatureGenerator()

    predictions = []
    warnings = []
    date_range = (request.date, (start_date + timedelta(days=request.days - 1)).strftime("%Y-%m-%d"))

    for day_offset in range(request.days):
        current_date = (start_date + timedelta(days=day_offset)).strftime("%Y-%m-%d")

        try:
            features_df = generator.generate_features(current_date)

            if request.district:
                features_df = features_df[
                    features_df['district'] == request.district
                ]

            for _, row in features_df.iterrows():
                risk_1d = float(row.get('risk_1day', 0.5))  # type: ignore[arg-type]
                risk_3d = float(row.get('risk_3day', 0.5))  # type: ignore[arg-type]
                risk_7d = float(row.get('risk_7day', 0.5))  # type: ignore[arg-type]

                if risk_1d >= 0.8:
                    risk_level = "high"
                elif risk_1d >= 0.6:
                    risk_level = "medium_high"
                elif risk_1d >= 0.4:
                    risk_level = "medium"
                elif risk_1d >= 0.2:
                    risk_level = "medium_low"
                else:
                    risk_level = "low"

                predictions.append(GridPrediction(
                    grid_id=row['grid_id'],
                    latitude=row.get('center_lat', 0),
                    longitude=row.get('center_lon', 0),
                    risk_1day=risk_1d,
                    risk_3day=risk_3d,
                    risk_7day=risk_7d,
                    risk_level=risk_level,
                    confidence=0.85,
                ))

        except Exception as e:
            logging.getLogger("cbpoa.grid").warning("Failed to generate features for %s: %s", current_date, e)
            warnings.append(f"Failed to generate features for {current_date}: {e}")
            continue

        if len(predictions) >= 50000:
            break

    return MultiDayPredictionResponse(
        predictions=predictions[:50000],
        total_grids=len(predictions),
        date_range=date_range,
        model_version="1.3.7",
        timestamp=datetime.now().isoformat(),
        partial=len(warnings) > 0,
        warnings=warnings,
    )


def _compute_grid_history(grid_id: str, days: int) -> dict:
    """Heavy synchronous parquet reads + per-row loop (called in thread pool)."""
    import pandas as pd

    district_map = _load_parquet(PROJECT_ROOT / "processed" / "grid_district_mapping.parquet")
    grid_info = district_map[district_map['grid_id'] == grid_id]

    if len(grid_info) == 0:
        raise HTTPException(status_code=404, detail="Grid not found")

    district = grid_info.iloc[0]['district_name']

    cases_df = load_cases_by_district_daily()
    cases_df['date'] = pd.to_datetime(cases_df['date'])

    end_date = datetime.now()
    start_date = end_date - timedelta(days=days)

    filtered = cases_df[
        (cases_df['date'] >= start_date) &
        (cases_df['date'] <= end_date) &
        (cases_df['district'] == district)
    ]

    history = []
    for _, row in filtered.iterrows():
        history.append({
            "date": row['date'].strftime("%Y-%m-%d"),
            "cases": int(row['total_cases']),
            "outpatient": int(row['outpatient_count']),
            "inpatient": int(row['inpatient_count']),
        })

    return {
        "grid_id": grid_id,
        "district": district,
        "history": history,
        "timestamp": datetime.now().isoformat(),
    }


@router.get("/grids/{grid_id}/history")
async def get_grid_history(
    grid_id: str,
    days: int = Query(30, ge=1, le=365, description="Number of days of history"),
):
    """
    Get historical data for a specific grid cell.

    Parquet reads + aggregation run in a thread pool to avoid blocking the
    async event loop.
    """
    return await asyncio.to_thread(_compute_grid_history, grid_id, days)