CA/backend/routers/cases.py

"""
医疗病例数据 API 路由

提供门诊和住院数据的统计、趋势、区域分布等接口
"""

from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel
from typing import Optional
from datetime import datetime, date
import asyncio
import pandas as pd
import json

from data.case_loader import load_data, get_combined_data, get_outpatient_data, get_inpatient_data, get_diagnoses, WUHAN_DISTRICTS, DATE_PATTERN

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


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

class StatsResponse(BaseModel):
    """统计数据响应"""
    total_outpatient: int
    total_inpatient: int
    date_range: dict
    top_districts: list
    top_diagnoses: list


class TrendPoint(BaseModel):
    """趋势数据点"""
    date: str
    outpatient: int
    inpatient: int
    total: int


class TrendResponse(BaseModel):
    """趋势数据响应"""
    trend: list[TrendPoint]
    summary: dict


class DistrictData(BaseModel):
    """区域数据"""
    district: str
    outpatient: int
    inpatient: int
    total: int
    outpatient_ratio: float
    inpatient_ratio: float


class DistrictsResponse(BaseModel):
    """区域分布响应"""
    districts: list[DistrictData]
    total: int


class RealtimeData(BaseModel):
    """实时数据"""
    today_outpatient: int
    today_inpatient: int
    today_total: int
    last_7d_avg: int
    change_ratio: float
    status: str


# ============== API Endpoints ==============

@router.get("/stats", response_model=StatsResponse, summary="获取病例统计数据")
async def get_cases_stats(
    diagnosis: Optional[str] = Query(None, description="Filter to single disease stats"),
):
    """
    获取病例总体统计信息

    - 总门诊量、总住院量
    - 数据日期范围
    - 就诊量前 10 的区域
    - 最常见诊断前 10
    """
    load_data()

    df_out = get_outpatient_data()
    df_in = get_inpatient_data()

    # 诊断过滤
    if diagnosis:
        df_out = df_out[df_out['初诊'].str.contains(diagnosis, na=False, case=False)]
        df_in = df_in[df_in['诊断名称'].str.contains(diagnosis, na=False, case=False)]

    # 计算统计
    total_outpatient = len(df_out)
    total_inpatient = len(df_in)

    # 日期范围（过滤后可能为空，需防御 NaT）
    all_dates = pd.concat([df_out['date'], df_in['date']]).dropna()
    date_start = all_dates.min().strftime("%Y-%m-%d") if len(all_dates) else ""
    date_end = all_dates.max().strftime("%Y-%m-%d") if len(all_dates) else ""

    # 区域统计
    out_districts = df_out[df_out['district'] != '未知']['district'].value_counts().head(10)
    in_districts = df_in[df_in['district'] != '其他']['district'].value_counts().head(10)

    combined_districts = pd.concat([out_districts, in_districts]).groupby(level=0).sum().nlargest(10)
    top_districts = [{"district": d, "count": int(c)} for d, c in combined_districts.items()]

    # 诊断统计
    out_diagnoses = df_out['初诊'].value_counts().head(10)
    in_diagnoses = df_in['诊断名称'].value_counts().head(10)

    top_diagnoses = [
        {"diagnosis": str(d), "outpatient": int(out_diagnoses.get(d, 0)), "inpatient": int(in_diagnoses.get(d, 0))}
        for d in set(list(out_diagnoses.index[:5]) + list(in_diagnoses.index[:5]))
    ][:10]

    return StatsResponse(
        total_outpatient=total_outpatient,
        total_inpatient=total_inpatient,
        date_range={
            "start": date_start,
            "end": date_end
        },
        top_districts=top_districts,
        top_diagnoses=top_diagnoses
    )


@router.get("/trend", response_model=TrendResponse, summary="获取病例趋势数据")
async def get_cases_trend(
    start_date: Optional[str] = Query(None, description="开始日期 (YYYY-MM-DD)"),
    end_date: Optional[str] = Query(None, description="结束日期 (YYYY-MM-DD)"),
    group_by: str = Query("day", description="分组粒度：day, week, month"),
    diagnosis: Optional[str] = Query(None, description="Filter by diagnosis name"),
):
    """
    获取病例时间趋势数据

    - 支持按日、周、月分组
    - 可指定日期范围
    - 返回门诊、住院、总计趋势
    """
    if start_date and not DATE_PATTERN.match(start_date):
        raise HTTPException(status_code=400, detail="Invalid start_date format. Use YYYY-MM-DD")
    if end_date and not DATE_PATTERN.match(end_date):
        raise HTTPException(status_code=400, detail="Invalid end_date format. Use YYYY-MM-DD")

    df = get_combined_data().copy()

    # 日期过滤
    if start_date:
        df = df[df['date'] >= pd.to_datetime(start_date)]
    if end_date:
        df = df[df['date'] <= pd.to_datetime(end_date)]

    # 诊断过滤
    if diagnosis:
        df = df[df['diagnosis'].str.contains(diagnosis, na=False, case=False)]

    # 分组
    if group_by == "week":
        df['period'] = df['date'].dt.to_period('W').dt.start_time
    elif group_by == "month":
        df['period'] = df['date'].dt.to_period('M').dt.start_time
    else:
        df['period'] = df['date'].dt.date

    # 聚合
    out_trend = df[df['type'] == 'outpatient'].groupby('period').size()
    in_trend = df[df['type'] == 'inpatient'].groupby('period').size()

    periods = sorted(set(out_trend.index.tolist() + in_trend.index.tolist()))

    trend = []
    total_out = total_in = 0
    for p in periods:
        out_count = int(out_trend.get(p, 0))
        in_count = int(in_trend.get(p, 0))
        total_out += out_count
        total_in += in_count
        trend.append(TrendPoint(
            date=str(p).split(' ')[0] if hasattr(p, 'strftime') else str(p)[:10],
            outpatient=out_count,
            inpatient=in_count,
            total=out_count + in_count
        ))

    return TrendResponse(
        trend=trend,
        summary={
            "total_outpatient": total_out,
            "total_inpatient": total_in,
            "period_count": len(periods),
            "avg_daily_outpatient": round(total_out / max(len(periods), 1), 2),
            "avg_daily_inpatient": round(total_in / max(len(periods), 1), 2),
        }
    )


def _compute_cases_districts(
    case_type: Optional[str],
    min_count: int,
    diagnosis: Optional[str],
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
) -> DistrictsResponse:
    """Run the full pandas aggregation pipeline (called in thread pool)."""
    df = get_combined_data()

    # 诊断过滤
    if diagnosis:
        df = df[df['diagnosis'].str.contains(diagnosis, na=False, case=False)]

    # 日期过滤
    if start_date:
        df = df[df['date'] >= pd.to_datetime(start_date)]
    if end_date:
        df = df[df['date'] <= pd.to_datetime(end_date)]

    # 类型过滤
    if case_type == "outpatient":
        df = df[df['type'] == 'outpatient']
    elif case_type == "inpatient":
        df = df[df['type'] == 'inpatient']

    # 过滤未知区域
    df = df[(df['district'] != '未知') & (df['district'] != '其他')]

    # 聚合
    district_stats = df.groupby(['district', 'type']).size().unstack(fill_value=0)

    if 'outpatient' not in district_stats.columns:
        district_stats['outpatient'] = 0
    if 'inpatient' not in district_stats.columns:
        district_stats['inpatient'] = 0

    district_stats['total'] = district_stats['outpatient'] + district_stats['inpatient']

    # 过滤
    district_stats = district_stats[district_stats['total'] >= min_count]
    district_stats = district_stats.sort_values('total', ascending=False)

    total = int(district_stats['total'].sum())

    districts = []
    for district, row in district_stats.iterrows():
        districts.append(DistrictData(
            district=str(district),
            outpatient=int(row['outpatient']),
            inpatient=int(row['inpatient']),
            total=int(row['total']),
            outpatient_ratio=round(float(row['outpatient']) / float(row['total']) * 100, 2) if row['total'] > 0 else 0,
            inpatient_ratio=round(float(row['inpatient']) / float(row['total']) * 100, 2) if row['total'] > 0 else 0
        ))

    return DistrictsResponse(districts=districts, total=total)


@router.get("/districts", response_model=DistrictsResponse, summary="获取区域分布数据")
async def get_cases_districts(
    case_type: Optional[str] = Query(None, description="病例类型：outpatient, inpatient, all"),
    min_count: int = Query(10, description="最小病例数过滤"),
    diagnosis: Optional[str] = Query(None, description="Filter by diagnosis name"),
    start_date: Optional[str] = Query(None, description="开始日期 (YYYY-MM-DD)"),
    end_date: Optional[str] = Query(None, description="结束日期 (YYYY-MM-DD)"),
):
    """
    获取病例区域分布数据

    - 支持按病例类型筛选
    - 可设置最小病例数过滤
    - 支持日期范围过滤
    - 返回各区门诊、住院量及占比

    Pandas processing runs in a thread pool to avoid blocking the async event loop.
    """
    return await asyncio.to_thread(
        _compute_cases_districts, case_type, min_count, diagnosis, start_date, end_date
    )


@router.get("/realtime", response_model=RealtimeData, summary="获取实时数据")
async def get_cases_realtime():
    """
    获取实时病例数据

    - 今日就诊量
    - 近 7 日平均值
    - 变化率
    - 状态评估 (正常/偏高/偏低)
    """
    df = get_combined_data()

    today = pd.Timestamp.today().normalize()
    last_7d = today - pd.Timedelta(days=7)

    # 今日数据
    today_data = df[df['date'] >= today]
    today_total = len(today_data)
    today_out = len(today_data[today_data['type'] == 'outpatient'])
    today_in = len(today_data[today_data['type'] == 'inpatient'])

    # 近 7 日平均
    last_7d_data = df[(df['date'] >= last_7d) & (df['date'] < today)]
    last_7d_avg = round(len(last_7d_data) / 7, 2) if len(last_7d_data) > 0 else 0

    # 变化率
    if last_7d_avg > 0:
        change_ratio = round((today_total - last_7d_avg) / last_7d_avg * 100, 2)
    else:
        change_ratio = 0.0

    # 状态评估
    if change_ratio > 20:
        status = "偏高"
    elif change_ratio < -20:
        status = "偏低"
    else:
        status = "正常"

    return RealtimeData(
        today_outpatient=today_out,
        today_inpatient=today_in,
        today_total=today_total,
        last_7d_avg=int(last_7d_avg),
        change_ratio=change_ratio,
        status=status
    )


class DiagnosesResponse(BaseModel):
    """诊断列表响应"""
    diagnoses: list[str]


@router.get("/diagnoses", response_model=DiagnosesResponse, summary="获取所有诊断名称列表")
async def get_diagnoses_list():
    """Returns deduplicated, sorted list of unique diagnosis names (cached, fast)."""
    diagnoses = get_diagnoses()
    return DiagnosesResponse(diagnoses=diagnoses)


# ============== Seasonal & Distribution Endpoints ==============

class SeasonalPoint(BaseModel):
    """月度聚合数据点"""
    month: int  # 1-12
    month_label: str  # "1月", "2月", ...
    outpatient: int
    inpatient: int
    total: int


class SeasonalResponse(BaseModel):
    """月度季节性响应"""
    monthly: list[SeasonalPoint]
    period_years: list[int]  # e.g. [2022, 2023, 2024]
    total_cases: int


class DiagnosisDistributionItem(BaseModel):
    """诊断分布数据项"""
    diagnosis: str
    outpatient: int
    inpatient: int
    total: int
    percentage: float


class DiagnosisDistributionResponse(BaseModel):
    """诊断分布响应"""
    diagnoses: list[DiagnosisDistributionItem]
    total_cases: int


# ============== Demographics Models ==============

class AgeBin(BaseModel):
    """年龄分段数据"""
    age_bin: int  # 0-17
    outpatient: int
    inpatient: int


class GenderSplit(BaseModel):
    """性别拆分数据"""
    outpatient: int
    inpatient: int


class GenderSplitData(BaseModel):
    """性别分布响应内层"""
    male: GenderSplit
    female: GenderSplit


class AgeDiagnosisMatrixItem(BaseModel):
    """年龄-诊断矩阵项"""
    age_group: str  # "0-1", "1-3", "3-6", "6-12", "12-18"
    diagnosis: str
    outpatient: int
    inpatient: int
    total: int


class DemographicsResponse(BaseModel):
    """人口统计响应"""
    age_distribution: list[AgeBin]
    gender_split: GenderSplitData
    age_diagnosis_matrix: list[AgeDiagnosisMatrixItem]


# ============== Disease Seasonality Models ==============

class DiseaseSeasonalityPoint(BaseModel):
    """疾病月度季节性数据点"""
    diagnosis: str
    month: int  # 1-12
    month_label: str  # "1月"-"12月"
    outpatient: int
    inpatient: int
    total: int


class DiseaseDistrictItem(BaseModel):
    """单个诊断的区域分布（按病例数排序的前若干区）"""
    diagnosis: str
    district: str
    total: int


class DiseaseSeasonalityResponse(BaseModel):
    """疾病季节性响应"""
    seasonality: list[DiseaseSeasonalityPoint]
    diagnoses: list[str]
    # 每个诊断的真实区域分布（按区聚合），使前端可为每个诊断显示其各自的"主要区域"
    diagnosis_districts: list[DiseaseDistrictItem]


@router.get("/seasonal", response_model=SeasonalResponse, summary="获取季节性月度聚合数据")
async def get_cases_seasonal(
    diagnosis: Optional[str] = Query(None, description="Filter by diagnosis name"),
):
    """
    按月聚合所有年份的病例数据

    - 返回 1-12 月各月门诊/住院/总计均值
    - 支持诊断过滤
    - 用于季节性分解图表
    """
    df = get_combined_data()

    if diagnosis:
        df = df[df['diagnosis'].str.contains(diagnosis, na=False, case=False)]

    # Extract month and aggregate
    df = df.copy()
    df['month'] = df['date'].dt.month
    years = sorted(df['date'].dt.year.unique().tolist())

    out_monthly = df[df['type'] == 'outpatient'].groupby('month').size()
    in_monthly = df[df['type'] == 'inpatient'].groupby('month').size()

    month_labels = ['1月', '2月', '3月', '4月', '5月', '6月',
                    '7月', '8月', '9月', '10月', '11月', '12月']

    monthly = []
    total_cases = 0
    for m in range(1, 13):
        out_count = int(out_monthly.get(m, 0))
        in_count = int(in_monthly.get(m, 0))
        total_cases += out_count + in_count
        monthly.append(SeasonalPoint(
            month=m,
            month_label=month_labels[m - 1],
            outpatient=out_count,
            inpatient=in_count,
            total=out_count + in_count,
        ))

    return SeasonalResponse(
        monthly=monthly,
        period_years=years,
        total_cases=total_cases,
    )


@router.get("/diagnosis-distribution", response_model=DiagnosisDistributionResponse, summary="获取诊断分布统计")
async def get_diagnosis_distribution(
    limit: int = Query(default=20, ge=1, le=50, description="Maximum diagnoses to return"),
):
    """
    获取诊断名称分布统计（门诊+住院分列）

    - 返回前 N 个诊断及门诊/住院/总计/占比
    - 用于诊断分布饼图、树图等
    """
    df = get_combined_data()

    # Compute O/I counts per diagnosis
    breakdown = df.groupby(['diagnosis', 'type']).size().unstack(fill_value=0)
    if 'outpatient' not in breakdown.columns:
        breakdown['outpatient'] = 0
    if 'inpatient' not in breakdown.columns:
        breakdown['inpatient'] = 0
    breakdown['total'] = breakdown['outpatient'] + breakdown['inpatient']
    breakdown = breakdown.sort_values('total', ascending=False).head(limit)

    grand_total = int(breakdown['total'].sum())

    diagnoses = []
    for diagnosis_name, row in breakdown.iterrows():
        diagnoses.append(DiagnosisDistributionItem(
            diagnosis=str(diagnosis_name),
            outpatient=int(row['outpatient']),
            inpatient=int(row['inpatient']),
            total=int(row['total']),
            percentage=round(float(row['total']) / float(grand_total) * 100, 2) if grand_total > 0 else 0,
        ))

    return DiagnosisDistributionResponse(
        diagnoses=diagnoses,
        total_cases=grand_total,
    )


# ============== Demographics Endpoint ==============

@router.get("/demographics", response_model=DemographicsResponse, summary="获取人口统计信息")
async def get_cases_demographics():
    """
    获取病例人口统计信息

    - 年龄分布（0-17岁，按1岁分段，仅住院数据）
    - 性别分布（仅住院数据）
    - 年龄-诊断矩阵（按年龄段分组，仅住院数据）

    注意：门诊数据不包含人口统计信息（性别/年龄），因此门诊计数均为 0。
    """
    df = get_inpatient_data()
    df = df.copy()
    df['age_bin'] = df['年龄'].clip(0, 17).astype(int)

    # --- Age distribution: 1-year bins from 0 to 17 ---
    age_counts = df.groupby('age_bin').size()
    age_distribution = [
        AgeBin(age_bin=a, outpatient=0, inpatient=int(age_counts.get(a, 0)))
        for a in range(0, 18)
    ]

    # --- Gender split ---
    gender_counts = df['性别'].value_counts()
    gender_split = GenderSplitData(
        male=GenderSplit(outpatient=0, inpatient=int(gender_counts.get('男性', 0))),
        female=GenderSplit(outpatient=0, inpatient=int(gender_counts.get('女性', 0))),
    )

    # --- Age-diagnosis matrix ---
    age_bins = [
        (0, 1, "0-1"), (1, 3, "1-3"), (3, 6, "3-6"),
        (6, 12, "6-12"), (12, 18, "12-18"),
    ]

    matrix_rows: list[AgeDiagnosisMatrixItem] = []
    for low, high, label in age_bins:
        group = df[(df['年龄'] >= low) & (df['年龄'] < high)]
        for diag, count in group['诊断名称'].value_counts().items():
            matrix_rows.append(AgeDiagnosisMatrixItem(
                age_group=label, diagnosis=str(diag),
                outpatient=0, inpatient=int(count), total=int(count),
            ))

    return DemographicsResponse(
        age_distribution=age_distribution,
        gender_split=gender_split,
        age_diagnosis_matrix=matrix_rows,
    )


# ============== Disease Seasonality Endpoint ==============

@router.get("/disease-seasonality", response_model=DiseaseSeasonalityResponse, summary="获取疾病季节性数据")
async def get_disease_seasonality(
    diagnosis: Optional[str] = Query(None, description="Filter by diagnosis name"),
):
    """
    获取各诊断的月度季节性分布数据

    - 基于门诊+住院合并数据
    - 按月聚合所有年份，返回 top 10 诊断的月度分布
    - 支持可选诊断过滤
    - 用于疾病季节性热力图、雷达图等
    """
    df = get_combined_data()

    if diagnosis:
        df = df[df['diagnosis'].str.contains(diagnosis, na=False, case=False)]

    # Extract month
    df = df.copy()
    df['month'] = df['date'].dt.month

    # Get top 10 diagnoses by total case count
    diag_totals = df.groupby('diagnosis').size().nlargest(10)
    top_diagnoses = diag_totals.index.tolist()

    month_labels = ['1月', '2月', '3月', '4月', '5月', '6月',
                    '7月', '8月', '9月', '10月', '11月', '12月']

    # Filter to top diagnoses
    df_top = df[df['diagnosis'].isin(top_diagnoses)]

    # Group by diagnosis + month
    breakdown = df_top.groupby(['diagnosis', 'month', 'type']).size().unstack(fill_value=0)
    if 'outpatient' not in breakdown.columns:
        breakdown['outpatient'] = 0
    if 'inpatient' not in breakdown.columns:
        breakdown['inpatient'] = 0

    seasonality: list[DiseaseSeasonalityPoint] = []
    for diag in top_diagnoses:
        for m in range(1, 13):
            row = breakdown.loc[(diag, m)] if (diag, m) in breakdown.index else None
            out_count = int(row['outpatient']) if row is not None else 0
            in_count = int(row['inpatient']) if row is not None else 0
            seasonality.append(DiseaseSeasonalityPoint(
                diagnosis=str(diag),
                month=m,
                month_label=month_labels[m - 1],
                outpatient=out_count,
                inpatient=in_count,
                total=out_count + in_count,
            ))

    # Per-diagnosis district distribution (real aggregation by diagnosis × district).
    # Previously the frontend showed the same "主要区域" for every diagnosis because
    # no per-diagnosis district data was exposed. Top 3 districts per diagnosis.
    df_districts = df_top[(df_top['district'] != '未知') & (df_top['district'] != '其他')]
    diag_district_counts = df_districts.groupby(['diagnosis', 'district']).size()
    diagnosis_districts: list[DiseaseDistrictItem] = []
    for diag in top_diagnoses:
        if diag not in diag_district_counts.index.get_level_values('diagnosis'):
            continue
        top_d = diag_district_counts.loc[diag].sort_values(ascending=False).head(3)
        for district_name, count in top_d.items():
            diagnosis_districts.append(DiseaseDistrictItem(
                diagnosis=str(diag),
                district=str(district_name),
                total=int(count),
            ))

    return DiseaseSeasonalityResponse(
        seasonality=seasonality,
        diagnoses=[str(d) for d in top_diagnoses],
        diagnosis_districts=diagnosis_districts,
    )