CA/backend/routers/analysis.py

"""
Router for CBPOA analysis endpoints
Time series trends, district aggregation, and weather-health correlations
"""
from fastapi import APIRouter, HTTPException, Query
from datetime import datetime, timedelta
from typing import List, Literal
import random

from pydantic import BaseModel, Field

from config import DATA_DIR, RISK_HIGH
from utils.date_helpers import get_latest_date
from utils.geojson import parse_geojson_file, load_districts
from utils.geo import point_in_polygon
from utils.risk import calculate_trend

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


class TrendResponse(BaseModel):
    """Response for trend data"""
    dates: List[str] = Field(..., description="Date labels")
    values: List[float] = Field(..., description="Risk values")
    trend: Literal["up", "down", "stable"] = Field(..., description="Trend direction")


class DistrictRisk(BaseModel):
    """District-level risk aggregation"""
    name: str = Field(..., description="District name")
    avg_risk: float = Field(..., description="Average risk value")
    high_risk_count: int = Field(..., description="Count of high risk grids")
    total_grids: int = Field(..., description="Total grids in district")
    total_cases: int = Field(..., description="Estimated total cases")


class DistrictsResponse(BaseModel):
    """Response for districts aggregation"""
    districts: List[DistrictRisk] = Field(..., description="District risk data")
    timestamp: str = Field(..., description="Response timestamp")


class CorrelationFactor(BaseModel):
    """Correlation factor data"""
    factor: str = Field(..., description="Factor name")
    correlation: float = Field(..., description="Correlation coefficient (-1 to 1)")
    significance: Literal["high", "medium", "low"] = Field(..., description="Statistical significance")
    description: str = Field(..., description="Factor description")


class CorrelationsResponse(BaseModel):
    """Response for correlations"""
    correlations: List[CorrelationFactor] = Field(..., description="Correlation factors")
    timestamp: str = Field(..., description="Response timestamp")


@router.get("/trend", response_model=TrendResponse)
async def get_trend(days: int = Query(default=7, ge=1, le=30)):
    """
    Get time series trend data from ACTUAL historical observations

    Args:
        days: Number of days for trend (1-30)

    Returns:
        Trend data with dates, values, and trend direction
    """
    latest_date = get_latest_date()

    try:
        base_date = datetime.strptime(latest_date, "%Y%m%d")
    except ValueError:
        raise HTTPException(status_code=500, detail="Invalid date format in data files")

    dates = []
    values = []

    for i in range(days):
        date = base_date - timedelta(days=days - 1 - i)
        date_str = date.strftime("%Y%m%d")
        filepath = DATA_DIR / f"risk_{date_str}.geojson"

        if filepath.exists():
            grids = parse_geojson_file(filepath)
            if grids:
                avg_risk = sum(g["risk_value"] for g in grids) / len(grids)
                values.append(round(avg_risk, 4))
            else:
                values.append(0)
        else:
            values.append(0)
        dates.append(date.strftime("%Y-%m-%d"))

    # Filter out zero values
    valid_data = [(d, v) for d, v in zip(dates, values) if v > 0]
    if valid_data:
        dates, values = zip(*valid_data)
        dates, values = list(dates), list(values)

    trend_direction = calculate_trend(values)

    return TrendResponse(
        dates=dates,
        values=values,
        trend=trend_direction,
    )


@router.get("/districts", response_model=DistrictsResponse)
async def get_districts():
    """
    Get district-level risk aggregation

    Returns:
        District-level risk data with averages and counts
    """
    latest_date = get_latest_date()
    filepath = DATA_DIR / f"risk_{latest_date}.geojson"

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

    grids = parse_geojson_file(filepath)
    districts = load_districts()

    if not districts:
        # Fallback: return city-wide aggregation
        avg_risk = sum(g["risk_value"] for g in grids) / len(grids) if grids else 0
        high_risk_count = sum(1 for g in grids if g["risk_value"] >= RISK_HIGH)

        return DistrictsResponse(
            districts=[
                DistrictRisk(
                    name="武汉市",
                    avg_risk=round(avg_risk, 4),
                    high_risk_count=high_risk_count,
                    total_grids=len(grids),
                    total_cases=int(len(grids) * avg_risk * 0.1)  # Mock case rate
                )
            ],
            timestamp=datetime.now().isoformat()
        )

    # Aggregate grids by district using point-in-polygon
    district_data = {d["name"]: {"grids": [], "high_risk": 0} for d in districts}
    unassigned = {"grids": [], "high_risk": 0}

    for grid in grids:
        assigned = False
        for district in districts:
            if point_in_polygon(grid["latitude"], grid["longitude"], district["coordinates"]):
                district_data[district["name"]]["grids"].append(grid)
                if grid["risk_value"] >= RISK_HIGH:
                    district_data[district["name"]]["high_risk"] += 1
                assigned = True
                break

        if not assigned:
            unassigned["grids"].append(grid)
            if grid["risk_value"] >= RISK_HIGH:
                unassigned["high_risk"] += 1

    # Build response
    result = []
    for district in districts:
        name = district["name"]
        grids_in_district = district_data[name]["grids"]

        if not grids_in_district:
            continue

        avg_risk = sum(g["risk_value"] for g in grids_in_district) / len(grids_in_district)
        high_risk_count = district_data[name]["high_risk"]

        # Mock total cases based on risk and grid count
        total_cases = int(len(grids_in_district) * avg_risk * 0.1)

        result.append(
            DistrictRisk(
                name=name,
                avg_risk=round(avg_risk, 4),
                high_risk_count=high_risk_count,
                total_grids=len(grids_in_district),
                total_cases=total_cases
            )
        )

    # Add unassigned as "其他" if significant
    if unassigned["grids"]:
        avg_risk = sum(g["risk_value"] for g in unassigned["grids"]) / len(unassigned["grids"])
        result.append(
            DistrictRisk(
                name="其他",
                avg_risk=round(avg_risk, 4),
                high_risk_count=unassigned["high_risk"],
                total_grids=len(unassigned["grids"]),
                total_cases=int(len(unassigned["grids"]) * avg_risk * 0.1)
            )
        )

    return DistrictsResponse(
        districts=result,
        timestamp=datetime.now().isoformat()
    )


@router.get("/correlations", response_model=CorrelationsResponse)
async def get_correlations():
    """
    Get weather-health correlation analysis

    Returns:
        Correlation factors with coefficients and significance
    """
    latest_date = get_latest_date()
    filepath = DATA_DIR / f"risk_{latest_date}.geojson"

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

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

    # Calculate mock correlations based on risk patterns
    # In production, this would use actual weather and health data
    avg_risk = sum(g["risk_value"] for g in grids) / len(grids)
    risk_variance = sum((g["risk_value"] - avg_risk) ** 2 for g in grids) / len(grids)

    # Generate realistic correlation coefficients
    correlations = [
        CorrelationFactor(
            factor="temperature",
            correlation=round(-0.45 - 0.1 * (avg_risk - 0.5), 3),
            significance="high" if risk_variance > 0.05 else "medium",
            description="Temperature vs risk: Lower temps correlate with higher risk"
        ),
        CorrelationFactor(
            factor="humidity",
            correlation=round(0.32 + 0.15 * (avg_risk - 0.5), 3),
            significance="medium",
            description="Humidity vs risk: Higher humidity slightly increases risk"
        ),
        CorrelationFactor(
            factor="PM2.5",
            correlation=round(0.58 + 0.1 * (avg_risk - 0.5), 3),
            significance="high",
            description="PM2.5 vs risk: Strong positive correlation"
        ),
        CorrelationFactor(
            factor="PM10",
            correlation=round(0.51 + 0.08 * (avg_risk - 0.5), 3),
            significance="high",
            description="PM10 vs risk: Moderate positive correlation"
        ),
        CorrelationFactor(
            factor="wind_speed",
            correlation=round(-0.28 - 0.05 * (avg_risk - 0.5), 3),
            significance="low",
            description="Wind speed vs risk: Higher wind disperses pollutants"
        ),
        CorrelationFactor(
            factor="population_density",
            correlation=round(0.42 + 0.12 * (avg_risk - 0.5), 3),
            significance="high",
            description="Population density vs risk: Dense areas show higher transmission"
        ),
    ]

    return CorrelationsResponse(
        correlations=correlations,
        timestamp=datetime.now().isoformat()
    )