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.
"""

import re
import pandas as pd
from pathlib import Path
from datetime import datetime

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

# Data cache
_cache = {
    "outpatient": None,
    "inpatient": None,
    "loaded_at": None,
}

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

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


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_data():
    """Load and cache outpatient + inpatient data from Excel files"""
    if _cache["loaded_at"] is not None:
        return

    try:
        # Load outpatient data
        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

        # Load inpatient data
        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()
    except Exception as e:
        raise RuntimeError(f"Data loading failed: {str(e)}")


def get_combined_data():
    """Return merged outpatient + inpatient data with unified diagnosis column"""
    load_data()

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

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

    return pd.concat([df_out, df_in], ignore_index=True)


def get_outpatient_data():
    """Return the cached outpatient dataframe"""
    load_data()
    return _cache["outpatient"]


def get_inpatient_data():
    """Return the cached inpatient dataframe"""
    load_data()
    return _cache["inpatient"]