CA/backend/routers/grid.py

from fastapi import APIRouter, HTTPException, Query
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional
import logging
import sys
import math

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

from models import (
    DistrictAggregation,
    HistoricalAggregationRequest,
    HistoricalAggregationResponse,
    GridGeoJSONResponse,
    GridPrediction,
    MultiDayPredictionRequest,
    MultiDayPredictionResponse,
)

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


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

    import pandas as pd

    cases_df = pd.read_parquet(PROJECT_ROOT / "processed" / "cases_by_district_daily.parquet")
    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['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['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')

    weather_df = pd.read_parquet(PROJECT_ROOT / "processed" / "weather" / "station_daily_2022.parquet")
    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=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 pd.notna(row['AQI']) else 0.0,
            avg_PM25=float(row['PM25']) if pd.notna(row['PM25']) else 0.0,
            avg_PM10=float(row['PM10']) if pd.notna(row['PM10']) else 0.0,
        ))

    return HistoricalAggregationResponse(
        aggregations=aggregations,
        total_records=len(aggregations),
        date_range=(start_date, end_date),
        timestamp=datetime.now().isoformat(),
    )


@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.
    """
    import pandas as pd

    try:
        grid_df = pd.read_parquet(PROJECT_ROOT / "processed" / "grid_100m_index.parquet")
    except FileNotFoundError:
        return GridGeoJSONResponse(type="FeatureCollection", features=[], timestamp=datetime.now().isoformat())

    try:
        district_map = pd.read_parquet(PROJECT_ROOT / "processed" / "grid_district_mapping.parquet")
    except FileNotFoundError:
        return GridGeoJSONResponse(type="FeatureCollection", features=[], timestamp=datetime.now().isoformat())

    merged = grid_df.merge(district_map, on='grid_id', how='left')

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

    try:
        cases_df = pd.read_parquet(PROJECT_ROOT / "processed" / "cases_by_district_daily.parquet")
    except FileNotFoundError:
        return GridGeoJSONResponse(type="FeatureCollection", features=[], timestamp=datetime.now().isoformat())
    cases_df['date'] = pd.to_datetime(cases_df['date']).dt.strftime('%Y-%m-%d')

    cases_df = cases_df[cases_df['date'] == date]

    merged = merged.merge(cases_df, left_on='district_name', right_on='district', how='left')
    merged['total_cases'] = merged['total_cases'].fillna(0).astype(int)

    def safe_float(val, default=0.0):
        try:
            v = float(val)
            return default if math.isnan(v) or math.isinf(v) else v
        except (TypeError, ValueError):
            return default

    def sanitize(obj):
        """Replace NaN/Inf with None for JSON serialization."""
        if isinstance(obj, float):
            if math.isnan(obj) or math.isinf(obj):
                return None
            return obj
        if isinstance(obj, dict):
            return {k: sanitize(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [sanitize(v) for v in obj]
        return obj

    features = []
    for _, row in merged.iterrows():
        lon = safe_float(row.get('center_lon'))
        lat = safe_float(row.get('center_lat'))
        if lon == 0.0 and lat == 0.0:
            continue

        # MVP: Simple risk calculation based on cases and population density
        total_cases = safe_float(row.get('total_cases', 0), 0)
        total_cases = int(total_cases)
        pop_density = safe_float(row.get('population_density', 0))

        # Risk formula: cases per 10k population + baseline
        risk_value = min(1.0, (total_cases / max(pop_density, 1)) * 10 + 0.1)
        
        if risk_value >= 0.7:
            risk_level = "high"
        elif risk_value >= 0.5:
            risk_level = "medium"
        elif risk_value >= 0.3:
            risk_level = "medium_low"
        else:
            risk_level = "low"
        
        district = row.get('district_name')
        if isinstance(district, float) and (math.isnan(district) or math.isinf(district)):
            district = "未知"

        feature = {
            "type": "Feature",
            "geometry": {
                "type": "Point",
                "coordinates": [lon, lat]
            },
            "properties": {
                "grid_id": str(row.get('grid_id', '')),
                "latitude": lat,
                "longitude": lon,
                "district": district,
                "total_cases": total_cases,
                "population_density": pop_density,
                "risk_value": round(risk_value, 3),
                "risk_level": risk_level,
            }
        }
        features.append(feature)

        if len(features) >= 10000:
            break

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


@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))
                risk_3d = float(row.get('risk_3day', 0.5))
                risk_7d = float(row.get('risk_7day', 0.5))

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


@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.
    """
    import pandas as pd

    district_map = pd.read_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 = pd.read_parquet(PROJECT_ROOT / "processed" / "cases_by_district_daily.parquet")
    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(),
    }