CA/backend/data/case_loader.py

"""
Shared data-loading module for case data (outpatient + inpatient).

Extracted from routers/cases.py so both cases and reports routers can use
the same cached data without circular imports.

Performance: reads from pre-generated Parquet files (~0.1s) instead of
Excel (~10s). Falls back to Excel if parquet files are missing.
"""

import re
import logging
import threading
from typing import Optional, cast
import pandas as pd
from pathlib import Path
from datetime import datetime

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

DATE_PATTERN = re.compile(r"^\d{4}-\d{2}-\d{2}$")

# Data cache
_cache: dict[str, Optional[pd.DataFrame | datetime]] = {
    "outpatient": None,
    "inpatient": None,
    "combined": None,
    "cases_by_district_daily": None,
    "loaded_at": None,
}

# Guards the lazy build so concurrent callers don't duplicate the load/concat.
_load_lock = threading.RLock()

# Canonical Wuhan administrative districts (13), matching the `name` field in
# Datas/武汉市.geojson. All district roll-ups must collapse to exactly these.
CANONICAL_DISTRICTS = [
    '江岸区', '江汉区', '硚口区', '汉阳区', '武昌区', '青山区', '洪山区',
    '东西湖区', '汉南区', '蔡甸区', '江夏区', '黄陂区', '新洲区',
]
# Bare (suffix-less) base name -> canonical 区-suffixed name.
_DISTRICT_BASE_TO_CANONICAL = {d[:-1]: d for d in CANONICAL_DISTRICTS}
_DISTRICT_SUFFIXES = ('区', '县', '市')


def normalize_district(name: str) -> str:
    """Map a district label to its canonical 区-suffixed form.

    The case parquet carries both bare ("武昌") and suffixed ("武昌区") spellings
    of the same district, which double-counts in any roll-up. This collapses
    them: known bare names map to their canonical form; already-suffixed names
    pass through unchanged; anything else gets a "区" appended.
    """
    if name is None:
        return name
    name = str(name).strip()
    if name in _DISTRICT_BASE_TO_CANONICAL:
        return _DISTRICT_BASE_TO_CANONICAL[name]
    if name.endswith(_DISTRICT_SUFFIXES):
        return name
    return f"{name}区"


# Wuhan district mapping
WUHAN_DISTRICTS = {
    '江岸区': ['江岸'],
    '江汉区': ['江汉'],
    '武昌区': ['武昌'],
    '洪山区': ['洪山'],
    '汉阳区': ['汉阳'],
    '东西湖区': ['东西湖'],
    '黄陂区': ['黄陂'],
    '硚口区': ['硚口'],
    '江夏区': ['江夏'],
    '青山区': ['青山'],
    '新洲区': ['新洲'],
    '蔡甸区': ['蔡甸'],
    '东湖新技术开发区': ['东湖新技术开发区', '光谷'],
    '经开（汉南）区': ['经开', '汉南', '经济开发区'],
    '东湖生态旅游风景区': ['东湖生态旅游风景区']
}

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


def _extract_district(addr: str) -> str:
    """Extract Wuhan district name from address string"""
    if pd.isna(addr):
        return '未知'
    addr = str(addr)
    for district, keywords in WUHAN_DISTRICTS.items():
        for kw in keywords:
            if kw in addr:
                return district
    return '其他'


def _load_from_parquet() -> bool:
    """Try to load data from pre-generated Parquet files. Returns True on success."""
    outpatient_path = PROCESSED_DIR / "cases_outpatient.parquet"
    inpatient_path = PROCESSED_DIR / "cases_inpatient.parquet"

    if not outpatient_path.exists() or not inpatient_path.exists():
        logger.info("Parquet files not found, falling back to Excel")
        return False

    try:
        _cache["outpatient"] = pd.read_parquet(outpatient_path)
        _cache["inpatient"] = pd.read_parquet(inpatient_path)
        _cache["loaded_at"] = datetime.now()
        logger.info("Loaded case data from Parquet (%d outpatient, %d inpatient)",
                    len(_cache["outpatient"]), len(_cache["inpatient"]))
        return True
    except Exception as e:
        logger.warning("Parquet load failed (%s), falling back to Excel", e)
        return False


def _load_from_excel():
    """Load data from Excel files (slow fallback)."""
    df_out = pd.read_excel(DATA_DIR / "view_门诊.xlsx")
    df_out['date'] = pd.to_datetime(df_out['门诊日期_re'])
    df_out['district'] = df_out['现住址区'].fillna('未知')
    _cache["outpatient"] = df_out

    df_in = pd.read_excel(DATA_DIR / "view_住院.xlsx")
    df_in['date'] = pd.to_datetime(df_in['入院日期_re'])
    df_in['district'] = df_in['现住址_脱敏'].apply(_extract_district)
    _cache["inpatient"] = df_in

    _cache["loaded_at"] = datetime.now()


def load_data():
    """Load and cache outpatient + inpatient data.

    Uses pre-generated Parquet files for fast loading (~0.1s).
    Falls back to Excel files (~10s) if Parquet is unavailable.
    """
    if _cache["loaded_at"] is not None:
        return

    with _load_lock:
        if _cache["loaded_at"] is not None:  # another thread loaded while we waited
            return
        if not _load_from_parquet():
            try:
                _load_from_excel()
            except Exception as e:
                raise RuntimeError(f"Data loading failed: {str(e)}")


def get_combined_data() -> pd.DataFrame:
    """Return merged outpatient + inpatient data with unified diagnosis column.

    Caches the result in memory after first call (~0.1s on cached hit).
    """
    if _cache["combined"] is not None:
        return cast(pd.DataFrame, _cache["combined"])

    with _load_lock:
        if _cache["combined"] is not None:  # built while we waited for the lock
            return cast(pd.DataFrame, _cache["combined"])

        load_data()

        df_out = cast(pd.DataFrame, _cache["outpatient"])
        df_in = cast(pd.DataFrame, _cache["inpatient"])

        df_out = df_out[['date', 'district', '初诊', '主诉']].copy()
        df_out['type'] = 'outpatient'
        df_out['diagnosis'] = df_out['初诊']

        df_in = df_in[['date', 'district', '诊断名称']].copy()
        df_in['type'] = 'inpatient'
        df_in['diagnosis'] = df_in['诊断名称']
        df_in['主诉'] = None

        _cache["combined"] = pd.concat([df_out, df_in], ignore_index=True)  # type: ignore[assignment]
        return cast(pd.DataFrame, _cache["combined"])


def get_diagnoses() -> list[str]:
    """Return sorted list of unique diagnosis names (fast: reads from cached DataFrames)."""
    load_data()
    out_diag = cast(pd.DataFrame, _cache["outpatient"])['初诊'].dropna().unique()
    in_diag = cast(pd.DataFrame, _cache["inpatient"])['诊断名称'].dropna().unique()
    return sorted(set(out_diag.tolist() + in_diag.tolist()))


def get_outpatient_data() -> pd.DataFrame:
    """Return the cached outpatient dataframe"""
    load_data()
    return _cache["outpatient"]  # type: ignore[return-value]


def get_inpatient_data() -> pd.DataFrame:
    """Return the cached inpatient dataframe"""
    load_data()
    return _cache["inpatient"]  # type: ignore[return-value]


def load_cases_by_district_daily() -> pd.DataFrame:
    """Load processed/cases_by_district_daily.parquet with districts normalized.

    The on-disk parquet carries both bare and 区-suffixed spellings of each
    district (26 labels = 13 districts × 2 spellings), so any groupby on the
    raw `district` column double-counts. This is the single data-access
    boundary: it normalizes labels to the canonical 13 and re-aggregates
    (sum of outpatient_count / inpatient_count / total_cases per
    normalized district + date), so every downstream consumer
    (analysis / grid / insights) sees clean, deduped 13-district data.

    Returns a copy with columns [date, district, outpatient_count,
    inpatient_count, total_cases]. Raises FileNotFoundError if the parquet
    is missing (callers handle this as they did before).
    """
    path = PROCESSED_DIR / "cases_by_district_daily.parquet"
    cached = _cache.get("cases_by_district_daily")
    if cached is not None:
        return cast(pd.DataFrame, cached).copy()

    with _load_lock:
        cached = _cache.get("cases_by_district_daily")
        if cached is not None:
            return cast(pd.DataFrame, cached).copy()

        df = pd.read_parquet(path)
        df["district"] = df["district"].map(normalize_district)
        agg = (
            df.groupby(["date", "district"], as_index=False)[
                ["outpatient_count", "inpatient_count", "total_cases"]
            ]
            .sum()
        )
        _cache["cases_by_district_daily"] = agg  # type: ignore[assignment]
        return agg.copy()