#!/usr/bin/env python3
"""
Medical Records ETL for Wuhan Respiratory Disease Risk Prediction Platform
Processes outpatient and inpatient records to daily district-level counts.
"""

import pandas as pd
from pathlib import Path

BASE_DIR = Path("/home/akiba/CA")
OUTPATIENT_SRC = BASE_DIR / "Datas/view_门诊.xlsx"
INPATIENT_SRC = BASE_DIR / "Datas/view_住院.xlsx"
OUT_DIR = BASE_DIR / "processed/medical"

RESPIRATORY_KEYWORDS = [
    "呼吸", "咳", "喘", "肺炎", "支气管", "咽痛", "感冒", "上呼吸道", "流感", "新冠"
]

RESPIRATORY_ICD_CODES = [f"J{i:02d}" for i in range(100)]


def is_respiratory_outpatient(chief_complaint: str) -> bool:
    if pd.isna(chief_complaint):
        return False
    return any(kw in str(chief_complaint) for kw in RESPIRATORY_KEYWORDS)


def is_respiratory_icd(code: str) -> bool:
    if pd.isna(code):
        return False
    code_str = str(code).strip().upper()
    if not code_str:
        return False
    base_code = code_str.split(".")[0]
    return base_code in RESPIRATORY_ICD_CODES


def extract_district(address: str) -> str:
    """Extract district name from address string."""
    if pd.isna(address):
        return ""
    address = str(address)
    wuhan_districts = [
        "江岸区", "江汉区", "硚口区", "汉阳区", "武昌区", "青山区", 
        "洪山区", "东西湖区", "汉南区", "蔡甸区", "江夏区", 
        "黄陂区", "新洲区", "东湖高新区", "武汉经开区"
    ]
    for district in wuhan_districts:
        if district in address:
            return district
    for district in ["江岸", "江汉", "硚口", "汉阳", "武昌", "青山", "洪山", 
                     "东西湖", "汉南", "蔡甸", "江夏", "黄陂", "新洲"]:
        if district in address:
            return district
    return ""


def process_outpatient():
    print("Loading outpatient data...")
    df = pd.read_excel(OUTPATIENT_SRC)
    print(f"  Total outpatient records: {len(df):,}")

    date_col = "门诊日期_re"
    district_col = "现住址区"
    complaint_col = "主诉"

    print("  Filtering respiratory cases...")
    df["is_respiratory"] = df[complaint_col].apply(is_respiratory_outpatient)
    df_resp = df[df["is_respiratory"]].copy()
    print(f"  Respiratory outpatient records: {len(df_resp):,}")

    df_resp["district"] = df_resp[district_col].apply(extract_district)
    df_filtered = df_resp[df_resp["district"] != ""].copy()
    print(f"  Records with valid Wuhan district: {len(df_filtered):,}")

    result = df_filtered.groupby([date_col, "district"]).size().reset_index(name="outpatient_count")
    result.columns = ["date", "district", "outpatient_count"]
    print(f"  Aggregated to {len(result):,} date-district combinations")

    return result


def process_inpatient():
    print("Loading inpatient data...")
    df = pd.read_excel(INPATIENT_SRC)
    print(f"  Total inpatient records: {len(df):,}")

    date_col = "入院日期_re"
    district_col = "现住址_脱敏"
    icd_col = "诊断编码"

    print("  Filtering respiratory cases (J00-J99)...")
    df["is_respiratory"] = df[icd_col].apply(is_respiratory_icd)
    df_resp = df[df["is_respiratory"]].copy()
    print(f"  Respiratory inpatient records: {len(df_resp):,}")

    df_resp["district"] = df_resp[district_col].apply(extract_district)
    df_filtered = df_resp[df_resp["district"] != ""].copy()
    print(f"  Records with valid Wuhan district: {len(df_filtered):,}")

    result = df_filtered.groupby([date_col, "district"]).size().reset_index(name="inpatient_count")
    result.columns = ["date", "district", "inpatient_count"]
    print(f"  Aggregated to {len(result):,} date-district combinations")

    return result


def main():
    print("=" * 60)
    print("Medical Records ETL - Wuhan Respiratory Disease Platform")
    print("=" * 60)

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    print("\n[1/2] Processing outpatient records...")
    outpatient_df = process_outpatient()
    outpatient_path = OUT_DIR / "outpatient_daily.parquet"
    outpatient_df.to_parquet(outpatient_path, index=False)
    print(f"  Saved: {outpatient_path}")
    print(f"  Records: {len(outpatient_df):,}, Cases: {outpatient_df['outpatient_count'].sum():,}")

    print("\n[2/2] Processing inpatient records...")
    inpatient_df = process_inpatient()
    inpatient_path = OUT_DIR / "inpatient_daily.parquet"
    inpatient_df.to_parquet(inpatient_path, index=False)
    print(f"  Saved: {inpatient_path}")
    print(f"  Records: {len(inpatient_df):,}, Cases: {inpatient_df['inpatient_count'].sum():,}")

    combined = outpatient_df.merge(inpatient_df, on=["date", "district"], how="outer").fillna(0)
    combined["outpatient_count"] = combined["outpatient_count"].astype(int)
    combined["inpatient_count"] = combined["inpatient_count"].astype(int)
    combined_path = OUT_DIR / "medical_daily.parquet"
    combined.to_parquet(combined_path, index=False)
    print(f"\n  Combined saved: {combined_path}")
    print(f"  Total date-district combinations: {len(combined):,}")

    print("\n" + "=" * 60)
    print("ETL Complete!")
    print("=" * 60)


if __name__ == "__main__":
    main()