CA/backend/routers/statistics.py

"""
统计分析 API 路由 (prefix /api/stats)

为前端统计仪表盘提供聚合后的临床、症状、发病率、环境相关性和时序数据。
所有数据从 processed/*.parquet 计算得出（文件型后端，无数据库）。

设计原则：
- 模块级缓存载入的 parquet（与其他路由一致）。
- 仅返回聚合结果，绝不直接 dump 原始行，保持 payload 小。
- 每个端点用 try/except 包裹，失败时返回合法的空结构（绝不让 UI 收到 500）。
- pandas 计算放入线程池 (asyncio.to_thread)，避免阻塞事件循环。
"""

import asyncio
import glob
import logging
import threading
from pathlib import Path
from typing import Optional, cast

import numpy as np
import pandas as pd
from fastapi import APIRouter, Query
from pydantic import BaseModel

from data.case_loader import (
    get_inpatient_data,
    get_outpatient_data,
    get_combined_data,
    load_cases_by_district_daily,
    normalize_district,
)

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

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

PROJECT_ROOT = Path(__file__).parent.parent.parent
PROCESSED_DIR = PROJECT_ROOT / "processed"

# ============== Module-level caches ==============

_cache: dict[str, object] = {}
_cache_lock = threading.RLock()

# 7 个污染物（与 feature snapshot 列名一致）
POLLUTANTS = ["AQI", "PM25", "PM10", "SO2", "NO2", "O3", "CO"]

# 主诉症状关键词（固定列表，子串匹配）。注意顺序：更具体的在前避免被宽泛词吞掉，
# 但因为是独立子串计数，顺序不影响结果，仅为可读性分组。
SYMPTOM_KEYWORDS = [
    "发热", "咳嗽", "咳", "喘息", "喘", "流涕", "鼻塞", "咽痛", "咽喉",
    "痰", "气促", "呼吸困难", "肺炎", "复诊", "随诊", "复查",
    "腹泻", "呕吐", "头痛", "乏力", "胸闷", "鼻涕", "发烧", "感冒",
]
REVISIT_KEYWORDS = ["复诊", "随诊", "复查"]


def _district_population() -> pd.Series:
    """各区人口（population_density 求和，按 grid_district_mapping 归属）。

    返回 index 为规范化区名(13)、值为人口的 Series。结果缓存。
    """
    cached = _cache.get("district_population")
    if cached is not None:
        return cast(pd.Series, cached)
    with _cache_lock:
        cached = _cache.get("district_population")
        if cached is not None:
            return cast(pd.Series, cached)
        mapping = pd.read_parquet(PROCESSED_DIR / "grid_district_mapping.parquet")
        grid = pd.read_parquet(
            PROCESSED_DIR / "grid_100m_with_dem_pop.parquet",
            columns=["grid_id", "population_density"],
        )
        joined = mapping.merge(grid, on="grid_id", how="inner")
        joined = joined.dropna(subset=["district_name"])
        joined["district_name"] = joined["district_name"].map(normalize_district)
        pop = joined.groupby("district_name")["population_density"].sum()
        _cache["district_population"] = pop
        return pop


def _feature_snapshots() -> pd.DataFrame:
    """合并所有可用的 features_*.parquet 快照（缓存）。

    用于污染物 vs 病例的相关性分析。每个快照按格点给出污染物 + 病例计数 + 区。
    """
    cached = _cache.get("features")
    if cached is not None:
        return cast(pd.DataFrame, cached)
    with _cache_lock:
        cached = _cache.get("features")
        if cached is not None:
            return cast(pd.DataFrame, cached)
        paths = sorted(glob.glob(str(PROCESSED_DIR / "features_*.parquet")))
        if not paths:
            df = pd.DataFrame(
                columns=POLLUTANTS + ["outpatient_count", "inpatient_count", "total_cases", "district"]
            )
        else:
            frames = [pd.read_parquet(p) for p in paths]
            df = pd.concat(frames, ignore_index=True)
        _cache["features"] = df
        return df


# ============== Response Models ==============


class KeyValueCount(BaseModel):
    bin_label: str
    count: int


class InpatientKpis(BaseModel):
    total_admissions: int
    median_los_days: float
    cure_rate: float
    emergency_admit_ratio: float


class LosByDisease(BaseModel):
    diagnosis: str
    p25: float
    median: float
    p75: float
    n: int


class LabelCount(BaseModel):
    outcome: Optional[str] = None
    route: Optional[str] = None
    count: int


class OutcomeCount(BaseModel):
    outcome: str
    count: int


class RouteCount(BaseModel):
    route: str
    count: int


class BmiByAge(BaseModel):
    age_band: str
    p25: float
    median: float
    p75: float
    n: int


class InpatientClinicalResponse(BaseModel):
    kpis: InpatientKpis
    los_histogram: list[KeyValueCount]
    los_by_disease: list[LosByDisease]
    outcome_counts: list[OutcomeCount]
    admission_route_counts: list[RouteCount]
    bmi_by_age_band: list[BmiByAge]


class SymptomItem(BaseModel):
    keyword: str
    count: int


class SymptomsResponse(BaseModel):
    symptoms: list[SymptomItem]
    revisit_ratio: float


class IncidenceItem(BaseModel):
    district: str
    total_cases: int
    population: float
    rate_per_10k: float


class IncidenceResponse(BaseModel):
    districts: list[IncidenceItem]


class CorrItem(BaseModel):
    pollutant: str
    corr_with_cases: float


class ScatterPoint(BaseModel):
    pm25: float
    aqi: float
    cases: float


class PairwiseCorr(BaseModel):
    a: str
    b: str
    corr: float


class EnvCorrelationResponse(BaseModel):
    correlation_matrix: list[CorrItem]
    scatter: list[ScatterPoint]
    pollutant_pairwise: list[PairwiseCorr]


class WeekdayPoint(BaseModel):
    weekday: str
    outpatient: int
    inpatient: int
    total: int


class MonthYearPoint(BaseModel):
    year: int
    month: int
    total: int


class YoYPoint(BaseModel):
    period: str
    current: int
    previous: int
    growth_pct: float


class TemporalResponse(BaseModel):
    weekday: list[WeekdayPoint]
    month_year: list[MonthYearPoint]
    yoy: list[YoYPoint]


# ============== Helpers ==============


def _empty_inpatient_clinical() -> InpatientClinicalResponse:
    return InpatientClinicalResponse(
        kpis=InpatientKpis(
            total_admissions=0, median_los_days=0.0,
            cure_rate=0.0, emergency_admit_ratio=0.0,
        ),
        los_histogram=[], los_by_disease=[], outcome_counts=[],
        admission_route_counts=[], bmi_by_age_band=[],
    )


def _safe_float(v) -> float:
    try:
        f = float(v)
        if np.isnan(f) or np.isinf(f):
            return 0.0
        return round(f, 4)
    except (TypeError, ValueError):
        return 0.0


# ============== Endpoint 1: inpatient clinical ==============


def _compute_inpatient_clinical() -> InpatientClinicalResponse:
    df = get_inpatient_data().copy()
    if df.empty:
        return _empty_inpatient_clinical()

    # LOS = (出院日期 - 入院日期).days, valid 0-60
    in_date = pd.to_datetime(df["入院日期"], errors="coerce")
    out_date = pd.to_datetime(df["出院日期"], errors="coerce")
    df["los"] = (out_date - in_date).dt.days
    df_los = df[(df["los"] >= 0) & (df["los"] <= 60)]

    total = len(df)
    median_los = float(df_los["los"].median()) if len(df_los) else 0.0

    outcome = df["出院情况"].fillna("未知")
    cure_n = int(outcome.isin(["治愈", "好转"]).sum())
    cure_rate = cure_n / total if total else 0.0

    route = df["入院途径"].fillna("未知")
    emerg_n = int((route == "急诊").sum())
    emerg_ratio = emerg_n / total if total else 0.0

    kpis = InpatientKpis(
        total_admissions=total,
        median_los_days=round(median_los, 2),
        cure_rate=round(cure_rate, 4),
        emergency_admit_ratio=round(emerg_ratio, 4),
    )

    # LOS histogram bins: 0,1,2,3,4,5,6,7,8-14,15+
    los_histogram: list[KeyValueCount] = []
    los_vals = df_los["los"]
    for b in range(0, 8):
        los_histogram.append(KeyValueCount(bin_label=str(b), count=int((los_vals == b).sum())))
    los_histogram.append(KeyValueCount(bin_label="8-14", count=int(((los_vals >= 8) & (los_vals <= 14)).sum())))
    los_histogram.append(KeyValueCount(bin_label="15+", count=int((los_vals >= 15).sum())))

    # LOS by disease (top 8 diagnoses by n)
    los_by_disease: list[LosByDisease] = []
    if len(df_los):
        top_diag = df_los["诊断名称"].value_counts().head(8).index.tolist()
        for d in top_diag:
            grp = df_los[df_los["诊断名称"] == d]["los"]
            los_by_disease.append(LosByDisease(
                diagnosis=str(d),
                p25=round(float(grp.quantile(0.25)), 2),
                median=round(float(grp.median()), 2),
                p75=round(float(grp.quantile(0.75)), 2),
                n=int(len(grp)),
            ))

    # outcome counts
    outcome_counts = [
        OutcomeCount(outcome=str(k), count=int(v))
        for k, v in outcome.value_counts().items()
    ]

    # admission route counts
    admission_route_counts = [
        RouteCount(route=str(k), count=int(v))
        for k, v in route.value_counts().items()
    ]

    # BMI by age band. BMI = 体重kg / (身高m)^2; plausible 8-40.
    bmi_by_age_band: list[BmiByAge] = []
    h = pd.to_numeric(df["身高"], errors="coerce")  # cm
    w = pd.to_numeric(df["体重"], errors="coerce")   # kg
    age = pd.to_numeric(df["年龄"], errors="coerce")
    bmi = w / ((h / 100.0) ** 2)
    bmi_df = pd.DataFrame({"age": age, "bmi": bmi})
    bmi_df = bmi_df[(bmi_df["bmi"] >= 8) & (bmi_df["bmi"] <= 40) & bmi_df["age"].notna()]
    age_bands = [(0, 3, "0-2"), (3, 6, "3-5"), (6, 9, "6-8"),
                 (9, 12, "9-11"), (12, 15, "12-14"), (15, 19, "15-18")]
    for lo, hi, label in age_bands:
        grp = bmi_df[(bmi_df["age"] >= lo) & (bmi_df["age"] < hi)]["bmi"]
        if len(grp) == 0:
            continue
        bmi_by_age_band.append(BmiByAge(
            age_band=label,
            p25=round(float(grp.quantile(0.25)), 2),
            median=round(float(grp.median()), 2),
            p75=round(float(grp.quantile(0.75)), 2),
            n=int(len(grp)),
        ))

    return InpatientClinicalResponse(
        kpis=kpis,
        los_histogram=los_histogram,
        los_by_disease=los_by_disease,
        outcome_counts=outcome_counts,
        admission_route_counts=admission_route_counts,
        bmi_by_age_band=bmi_by_age_band,
    )


@router.get("/inpatient-clinical", response_model=InpatientClinicalResponse, summary="住院临床统计")
async def inpatient_clinical():
    """住院临床概览：KPI、住院天数(LOS)分布、费用分布、转归、入院途径、BMI 分布。"""
    try:
        return await asyncio.to_thread(_compute_inpatient_clinical)
    except Exception:
        logger.exception("inpatient-clinical failed")
        return _empty_inpatient_clinical()


# ============== Endpoint 2: symptoms ==============


def _compute_symptoms(top: int) -> SymptomsResponse:
    df = get_outpatient_data()
    if df.empty or "主诉" not in df.columns:
        return SymptomsResponse(symptoms=[], revisit_ratio=0.0)
    chief = df["主诉"].dropna().astype(str)
    total = len(chief)
    if total == 0:
        return SymptomsResponse(symptoms=[], revisit_ratio=0.0)

    counts: list[SymptomItem] = []
    for kw in SYMPTOM_KEYWORDS:
        c = int(chief.str.contains(kw, regex=False).sum())
        if c > 0:
            counts.append(SymptomItem(keyword=kw, count=c))
    counts.sort(key=lambda x: x.count, reverse=True)
    counts = counts[:top]

    revisit_mask = chief.str.contains("|".join(REVISIT_KEYWORDS), regex=True)
    revisit_ratio = float(revisit_mask.sum()) / total if total else 0.0

    return SymptomsResponse(symptoms=counts, revisit_ratio=round(revisit_ratio, 4))


@router.get("/symptoms", response_model=SymptomsResponse, summary="门诊主诉症状词频")
async def symptoms(top: int = Query(20, ge=1, le=50, description="返回前 N 个症状词")):
    """从门诊主诉中提取固定症状关键词的出现频次，并计算复诊比例。"""
    try:
        return await asyncio.to_thread(_compute_symptoms, top)
    except Exception:
        logger.exception("symptoms failed")
        return SymptomsResponse(symptoms=[], revisit_ratio=0.0)


# ============== Endpoint 3: incidence rate ==============


def _compute_incidence() -> IncidenceResponse:
    daily = load_cases_by_district_daily()
    if daily.empty:
        return IncidenceResponse(districts=[])
    case_totals = daily.groupby("district")["total_cases"].sum()
    pop = _district_population()

    items: list[IncidenceItem] = []
    for district in case_totals.index:
        total_cases = int(case_totals.get(district, 0))
        population = float(pop.get(district, 0.0))
        rate = (total_cases / population * 10000) if population > 0 else 0.0
        items.append(IncidenceItem(
            district=str(district),
            total_cases=total_cases,
            population=round(population, 1),
            rate_per_10k=round(rate, 2),
        ))
    items.sort(key=lambda x: x.rate_per_10k, reverse=True)
    return IncidenceResponse(districts=items)


@router.get("/incidence-rate", response_model=IncidenceResponse, summary="各区发病率")
async def incidence_rate():
    """各区病例总数 / 区人口 * 10000，得到每万人发病率（13 区）。"""
    try:
        return await asyncio.to_thread(_compute_incidence)
    except Exception:
        logger.exception("incidence-rate failed")
        return IncidenceResponse(districts=[])


# ============== Endpoint 4: env correlation ==============


def _compute_env_correlation() -> EnvCorrelationResponse:
    df = _feature_snapshots()
    if df.empty or "total_cases" not in df.columns:
        return EnvCorrelationResponse(correlation_matrix=[], scatter=[], pollutant_pairwise=[])

    # 污染物 vs 病例 的 Pearson 相关（按格点，汇集所有快照）
    correlation_matrix: list[CorrItem] = []
    cases = pd.to_numeric(df["total_cases"], errors="coerce")
    for p in POLLUTANTS:
        if p not in df.columns:
            continue
        series = pd.to_numeric(df[p], errors="coerce")
        valid = series.notna() & cases.notna()
        if valid.sum() < 2 or series[valid].std() == 0 or cases[valid].std() == 0:
            corr = 0.0
        else:
            corr = float(series[valid].corr(cases[valid]))
        correlation_matrix.append(CorrItem(pollutant=p, corr_with_cases=_safe_float(corr)))

    # scatter: 采样 cases>0 的格点（up to 500）
    scatter: list[ScatterPoint] = []
    has_cols = all(c in df.columns for c in ["PM25", "AQI", "total_cases"])
    if has_cols:
        sdf = df[["PM25", "AQI", "total_cases"]].copy()
        sdf = sdf[pd.to_numeric(sdf["total_cases"], errors="coerce") > 0].dropna()
        if len(sdf) > 500:
            sdf = sdf.sample(n=500, random_state=42)
        for _, r in sdf.iterrows():
            scatter.append(ScatterPoint(
                pm25=_safe_float(r["PM25"]),
                aqi=_safe_float(r["AQI"]),
                cases=_safe_float(r["total_cases"]),
            ))

    # pollutant pairwise (upper triangle)
    pollutant_pairwise: list[PairwiseCorr] = []
    present = [p for p in POLLUTANTS if p in df.columns]
    pol_df = df[present].apply(pd.to_numeric, errors="coerce")
    corr_mat = pol_df.corr()
    for i, a in enumerate(present):
        for b in present[i + 1:]:
            try:
                v = corr_mat.loc[a, b]
            except KeyError:
                v = 0.0
            pollutant_pairwise.append(PairwiseCorr(a=a, b=b, corr=_safe_float(v)))

    return EnvCorrelationResponse(
        correlation_matrix=correlation_matrix,
        scatter=scatter,
        pollutant_pairwise=pollutant_pairwise,
    )


@router.get("/env-correlation", response_model=EnvCorrelationResponse, summary="环境-病例相关性")
async def env_correlation():
    """污染物与病例的相关矩阵、PM2.5/AQI 散点、污染物两两相关（热力图）。"""
    try:
        return await asyncio.to_thread(_compute_env_correlation)
    except Exception:
        logger.exception("env-correlation failed")
        return EnvCorrelationResponse(correlation_matrix=[], scatter=[], pollutant_pairwise=[])


# ============== Endpoint 5: temporal ==============

_WEEKDAY_LABELS = ["周一", "周二", "周三", "周四", "周五", "周六", "周日"]


def _compute_temporal() -> TemporalResponse:
    df = get_combined_data().copy()
    if df.empty:
        return TemporalResponse(weekday=[], month_year=[], yoy=[])
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df = df[df["date"].notna()]
    if df.empty:
        return TemporalResponse(weekday=[], month_year=[], yoy=[])

    # weekday (0=周一..6=周日)
    df["wd"] = df["date"].dt.weekday
    weekday: list[WeekdayPoint] = []
    for wd in range(7):
        sub = df[df["wd"] == wd]
        out = int((sub["type"] == "outpatient").sum())
        inp = int((sub["type"] == "inpatient").sum())
        weekday.append(WeekdayPoint(
            weekday=_WEEKDAY_LABELS[wd], outpatient=out, inpatient=inp, total=out + inp,
        ))

    # month_year (seasonality grid)
    df["year"] = df["date"].dt.year
    df["month"] = df["date"].dt.month
    my = df.groupby(["year", "month"]).size()
    month_year = [
        MonthYearPoint(year=int(y), month=int(m), total=int(c))
        for (y, m), c in my.items()
    ]
    month_year.sort(key=lambda x: (x.year, x.month))

    # yoy: monthly current vs same-month-prior-year (only if multiple years exist)
    yoy: list[YoYPoint] = []
    years = sorted(df["year"].unique().tolist())
    if len(years) > 1:
        monthly_totals = {(int(y), int(m)): int(c) for (y, m), c in my.items()}
        for (y, m), cur in sorted(monthly_totals.items()):
            prev = monthly_totals.get((y - 1, m))
            if prev is None:
                continue
            growth = ((cur - prev) / prev * 100) if prev else 0.0
            yoy.append(YoYPoint(
                period=f"{y}-{m:02d}",
                current=cur,
                previous=prev,
                growth_pct=round(growth, 2),
            ))

    return TemporalResponse(weekday=weekday, month_year=month_year, yoy=yoy)


@router.get("/temporal", response_model=TemporalResponse, summary="时序统计")
async def temporal():
    """按星期、年-月（季节性网格）聚合，以及同比(YoY)增长（若有多年数据）。"""
    try:
        return await asyncio.to_thread(_compute_temporal)
    except Exception:
        logger.exception("temporal failed")
        return TemporalResponse(weekday=[], month_year=[], yoy=[])