""" 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()