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,
    "loaded_at": None,
}

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

# 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]