step1_data_preprocessing.py

"""
NIPT 数据第一步处理：
1. 读取 data.xlsx；
2. 统一中英文列名；
3. 类型转换、日期解析、孕周解析、标签编码；
4. 生成 GC 质控标记 Quality_Pass；
5. 保存清洗后的数据。
"""

import argparse
import os
import re
from datetime import datetime

import numpy as np
import pandas as pd


RENAME_MAP = {
    "序号": "Sample_ID",
    "样本序号": "Sample_ID",
    "孕妇代码": "Patient_ID",
    "年龄": "Age",
    "孕妇年龄": "Age",
    "身高": "Height",
    "孕妇身高": "Height",
    "体重": "Weight",
    "孕妇体重": "Weight",
    "末次月经": "LMP_Date",
    "IVF妊娠": "IVF_Original",
    "IVF 妊娠方式": "IVF_Original",
    "检测日期": "Test_Date",
    "检测抽血次数": "Blood_Test_Number",
    "检测孕周": "Gestational_Week",
    "孕妇BMI": "BMI",
    "孕妇BMI指标": "BMI",
    "BMI": "BMI",
    "原始读段数": "Total_Reads",
    "原始测序数据的总读段数": "Total_Reads",
    "在参考基因组上比对的比例": "Mapped_Ratio",
    "总读段数中在参考基因组上比对的比例": "Mapped_Ratio",
    "重复读段的比例": "Duplicate_Ratio",
    "总读段数中重复读段的比例": "Duplicate_Ratio",
    "唯一比对的读段数": "Unique_Reads",
    "总读段数中唯一比对的读段数": "Unique_Reads",
    "GC含量": "GC_Total",
    "GC 含量": "GC_Total",
    "13号染色体的Z值": "Z_13",
    "18号染色体的Z值": "Z_18",
    "21号染色体的Z值": "Z_21",
    "X染色体的Z值": "Z_X",
    "Y染色体的Z值": "Z_Y",
    "Y染色体浓度": "Y_Conc",
    "X染色体浓度": "X_Conc",
    "13号染色体的GC含量": "GC_13",
    "18号染色体的GC含量": "GC_18",
    "21号染色体的GC含量": "GC_21",
    "被过滤掉读段数的比例": "Filtered_Reads_Ratio",
    "被过滤掉的读段数占总读段数的比例": "Filtered_Reads_Ratio",
    "染色体的非整倍体": "Aneuploidy_Status",
    "检测出的13号，18号，21号染色体非整倍体": "Aneuploidy_Status",
    "怀孕次数": "Pregnancy_Count",
    "生产次数": "Birth_Count",
    "胎儿是否健康": "Healthy_Outcome",
}


OUTPUT_COLUMNS = [
    "Sample_ID",
    "Patient_ID",
    "Fetus_Sex",
    "Original_Row",
    "Age",
    "Height",
    "Weight",
    "LMP_Date",
    "IVF_Original",
    "IVF",
    "Test_Date",
    "Blood_Test_Number",
    "Gestational_Week",
    "GA",
    "BMI",
    "Total_Reads",
    "Mapped_Ratio",
    "Duplicate_Ratio",
    "Unique_Reads",
    "GC_Total",
    "Z_13",
    "Z_18",
    "Z_21",
    "Z_X",
    "Z_Y",
    "Y_Conc",
    "X_Conc",
    "GC_13",
    "GC_18",
    "GC_21",
    "Filtered_Reads_Ratio",
    "Aneuploidy_Status",
    "Pregnancy_Count",
    "Pregnancy_Count_Original",
    "Birth_Count",
    "Healthy_Outcome",
    "Quality_Pass",
]


NUMERIC_COLUMNS = [
    "Age",
    "Height",
    "Weight",
    "Blood_Test_Number",
    "BMI",
    "Total_Reads",
    "Mapped_Ratio",
    "Duplicate_Ratio",
    "Unique_Reads",
    "GC_Total",
    "Z_13",
    "Z_18",
    "Z_21",
    "Z_X",
    "Z_Y",
    "Y_Conc",
    "X_Conc",
    "GC_13",
    "GC_18",
    "GC_21",
    "Filtered_Reads_Ratio",
    "Birth_Count",
]


def parse_gestational_week(week_str):
    """将 '11w+6'、'16W+1'、'12w'、'12' 等孕周文本解析为浮点周数。"""
    if pd.isna(week_str):
        return np.nan
    text = str(week_str).lower().strip()
    if not text:
        return np.nan

    match_wd = re.match(r"(\d+)\s*w\s*\+?\s*(\d+)\s*d?", text)
    if match_wd:
        return float(match_wd.group(1)) + float(match_wd.group(2)) / 7.0

    match_w = re.match(r"(\d+)\s*w", text)
    if match_w:
        return float(match_w.group(1))

    match_digits = re.match(r"(\d+)", text)
    if match_digits:
        return float(match_digits.group(1))

    return np.nan


def _parse_excel_like_date(value):
    """兼容 Excel 序列日期、YYYYMMDD 整数以及常见日期字符串。"""
    if pd.isna(value):
        return pd.NaT
    if isinstance(value, (datetime, pd.Timestamp)):
        return pd.Timestamp(value)

    text = str(value).strip()
    if text in ("", "nan", "NaT", "None"):
        return pd.NaT
    if text.endswith(".0") and text[:-2].isdigit():
        text = text[:-2]

    if text.isdigit():
        num = int(text)
        if 10000000 <= num <= 99999999:
            return pd.to_datetime(text, format="%Y%m%d")
        if 10000 <= num <= 60000:
            return pd.Timestamp("1899-12-30") + pd.Timedelta(days=num)

    for date_format in ("%Y-%m-%d", "%Y/%m/%d", "%Y.%m.%d", "%Y%m%d"):
        try:
            return pd.Timestamp(datetime.strptime(text, date_format))
        except ValueError:
            continue

    try:
        return pd.to_datetime(text)
    except (ValueError, TypeError):
        return pd.NaT


def _parse_date_series(series):
    return series.map(_parse_excel_like_date)


def _encode_ivf(value):
    if pd.isna(value):
        return np.nan
    return 1 if "IVF" in str(value).upper() else 0


def _encode_pregnancy_count(value):
    if pd.isna(value):
        return np.nan
    text = str(value).strip()
    if text in {"≥3", ">=3", "3+", ">2"}:
        return 3.0
    try:
        return float(text)
    except ValueError:
        return np.nan


def _encode_healthy(value):
    if pd.isna(value):
        return np.nan
    text = str(value).strip().lower()
    if text in {"是", "健康", "正常", "1", "true", "yes", "y"}:
        return 1
    if text in {"否", "异常", "不健康", "0", "false", "no", "n"}:
        return 0
    return np.nan


def _clean_text(value):
    if pd.isna(value):
        return np.nan
    text = str(value).strip()
    return text if text else np.nan


def _normalize_columns(df):
    """按别名映射统一列名，并丢弃空白/Unnamed 占位列。"""
    df = df.copy()
    rename_map = {}
    for col in df.columns:
        clean_name = str(col).strip()
        if clean_name in RENAME_MAP:
            rename_map[col] = RENAME_MAP[clean_name]
    df = df.rename(columns=rename_map)

    drop_cols = [
        col
        for col in df.columns
        if col is None
        or (isinstance(col, str) and (col.strip() == "" or col.startswith("Unnamed")))
    ]
    if drop_cols:
        df = df.drop(columns=drop_cols)
    return df


def _to_numeric(df, columns):
    for col in columns:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")


def load_nipt_data(xlsx_path="data.xlsx"):
    """用 pandas/openpyxl 读取 data.xlsx 中的男/女胎工作表，返回 (df_male, df_female)。"""
    if not os.path.exists(xlsx_path):
        raise FileNotFoundError(
            f"未找到数据源。请确保当前目录下存在 '{xlsx_path}'。"
        )

    print(f"正在从 Excel 附件 '{xlsx_path}' 中读取男/女胎检测数据...")
    try:
        excel_file = pd.ExcelFile(xlsx_path)
    except ImportError as exc:
        raise ImportError(
            "读取 Excel 需要 openpyxl，请先安装：pip install openpyxl"
        ) from exc

    male_sheet = next(
        (name for name in excel_file.sheet_names if "男" in name), None
    )
    female_sheet = next(
        (name for name in excel_file.sheet_names if "女" in name), None
    )
    if male_sheet is None or female_sheet is None:
        male_sheet, female_sheet = excel_file.sheet_names[:2]

    return excel_file.parse(male_sheet), excel_file.parse(female_sheet)


def preprocess_sheet(raw_df, fetus_sex, sample_start=1):
    """对单个工作表完成列名标准化、类型转换和基本清洗。"""
    df = _normalize_columns(raw_df)

    for col in OUTPUT_COLUMNS:
        if col not in df.columns:
            df[col] = np.nan

    df["Fetus_Sex"] = fetus_sex
    if "Sample_ID" in df.columns:
        df["Original_Row"] = pd.to_numeric(df["Sample_ID"], errors="coerce")
    else:
        df["Original_Row"] = np.nan
    df["Sample_ID"] = np.arange(sample_start, sample_start + len(df))

    df["LMP_Date"] = _parse_date_series(df["LMP_Date"])
    df["Test_Date"] = _parse_date_series(df["Test_Date"])
    _to_numeric(df, NUMERIC_COLUMNS)

    df["Pregnancy_Count_Original"] = df["Pregnancy_Count"].map(_clean_text)
    df["Pregnancy_Count"] = df["Pregnancy_Count"].map(_encode_pregnancy_count)
    df["Gestational_Week"] = df["Gestational_Week"].map(_clean_text)
    df["GA"] = df["Gestational_Week"].map(parse_gestational_week)
    df["IVF"] = df["IVF_Original"].map(_encode_ivf)
    df["Healthy_Outcome"] = df["Healthy_Outcome"].map(_encode_healthy)
    df["Aneuploidy_Status"] = df["Aneuploidy_Status"].map(_clean_text)

    df["Quality_Pass"] = np.where(
        df["GC_Total"].between(0.4, 0.6, inclusive="both"), 1, 0
    )
    return df[OUTPUT_COLUMNS]


def print_data_quality(df_male, df_female, df_all):
    """打印清洗后的数据质量摘要。"""
    print("\n" + "=" * 60)
    print("第一步数据处理完成，质量摘要：")
    print(f"男胎数据: {len(df_male)} 行, {df_male['Patient_ID'].nunique()} 位孕妇")
    print(
        f"女胎数据: {len(df_female)} 行, "
        f"{df_female['Patient_ID'].nunique()} 位孕妇"
    )
    print(f"合并后总记录数: {len(df_all)} 行")
    print(f"孕周解析失败: {int(df_all['GA'].isna().sum())} 行")
    print(
        f"GC 质控合格 (Quality_Pass=1): "
        f"{int(df_all['Quality_Pass'].sum())}/{len(df_all)}"
    )
    print(
        "健康标签分布: "
        f"{df_all['Healthy_Outcome'].value_counts(dropna=False).to_dict()}"
    )

    missing = df_all.isna().sum()
    missing = missing[missing > 0]
    if len(missing):
        print("非零缺失值统计:")
        for col, count in missing.items():
            print(f"  {col}: {count}")


def preprocess_nipt_data(
    xlsx_path="data.xlsx",
    save_path="nipt_processed_data.csv",
    save=True,
):
    """执行第一步数据处理，返回 (df_male, df_female, df_all)。"""
    raw_male, raw_female = load_nipt_data(xlsx_path=xlsx_path)

    df_male = preprocess_sheet(raw_male, fetus_sex="Male", sample_start=1)
    df_female = preprocess_sheet(
        raw_female, fetus_sex="Female", sample_start=len(df_male) + 1
    )
    df_all = pd.concat([df_male, df_female], ignore_index=True)

    print_data_quality(df_male, df_female, df_all)

    if save and save_path:
        df_all.to_csv(save_path, index=False, encoding="utf-8-sig")
        print(f"\n已保存预处理结果: {save_path}")

    return df_male, df_female, df_all


def main():
    parser = argparse.ArgumentParser(description="NIPT 数据第一步处理")
    parser.add_argument("--xlsx", default="data.xlsx", help="Excel 数据文件")
    parser.add_argument(
        "--output", default="nipt_processed_data.csv", help="清洗后输出 CSV"
    )
    parser.add_argument(
        "--no-save", action="store_true", help="不保存清洗后的 CSV"
    )
    args = parser.parse_args()

    preprocess_nipt_data(
        xlsx_path=args.xlsx,
        save_path=None if args.no_save else args.output,
    )


if __name__ == "__main__":
    main()