from __future__ import annotations

import argparse
import re
from pathlib import Path

import pandas as pd


FULL_COLUMNS = [
    "RID",
    "PTID",
    "label",
    "VISCODE",
    "EXAMDATE",
    "DX_bl",
    "DX",
    "AGE",
    "PTGENDER",
    "PTEDUCAT",
    "PTETHCAT",
    "PTRACCAT",
    "PTMARRY",
    "APOE4",
    "FDG",
    "PIB",
    "AV45",
    "ABETA",
    "TAU",
    "PTAU",
    "CDRSB",
    "ADAS11",
    "ADAS13",
    "ADASQ4",
    "MMSE",
    "RAVLT_immediate",
    "RAVLT_learning",
    "RAVLT_forgetting",
    "RAVLT_perc_forgetting",
    "LDELTOTAL",
    "DIGITSCOR",
    "TRABSCOR",
    "FAQ",
    "MOCA",
    "Years_bl",
]

CONVERSION_COLUMNS = [
    "PTID",
    "label",
    "DX.1",
    "mPACCdigit",
    "mPACCtrailsB",
    "Ventricles（心室）",
    "Hippocampus（海马）",
    "WholeBrain（全脑）",
    "Entorhinal（内嗅觉）",
    "Fusiform（梭形）",
    "MidTemp（中点温度）",
    "ICV",
]


def normalize_subject_id(value: object) -> str:
    if pd.isna(value):
        return ""
    text = str(value).strip().upper()
    text = text.replace("_S_", "S").replace("_", "")
    return re.sub(r"[^A-Z0-9]", "", text)


def clean_numeric(value: object) -> float | pd.NA:
    if pd.isna(value):
        return pd.NA
    text = str(value).strip()
    if not text:
        return pd.NA
    if text.startswith(">"):
        text = text[1:]
    try:
        return float(text)
    except ValueError:
        return pd.NA


def read_xlsx(path: Path, columns: list[str]) -> pd.DataFrame:
    df = pd.read_excel(path)
    available = [col for col in columns if col in df.columns]
    if "PTID" not in available:
        raise ValueError(f"{path} does not contain PTID.")
    df = df[available].copy()
    df["subject_id_norm"] = df["PTID"].map(normalize_subject_id)
    df = df[df["subject_id_norm"] != ""].copy()
    df = df.drop_duplicates("subject_id_norm", keep="first")
    return df


def add_numeric_clean_columns(df: pd.DataFrame) -> pd.DataFrame:
    for col in ["FDG", "PIB", "AV45", "ABETA", "TAU", "PTAU"]:
        if col in df.columns:
            df[f"{col}_num"] = df[col].map(clean_numeric)
    return df


def write_split_manifests(enriched: pd.DataFrame, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    for split_name, split_df in enriched.groupby("split", sort=False):
        split_df.to_csv(out_dir / f"{split_name}_clinical.csv", index=False)


def summarize(enriched: pd.DataFrame, out_md: Path, out_csv: Path) -> None:
    rows = []
    total = len(enriched)
    for col in [
        "clinical_label",
        "dx",
        "dx_bl",
        "conversion_label",
        "age",
        "sex",
        "apoe4",
        "mmse",
        "cdrsb",
        "adas11",
        "adas13",
        "faq",
        "moca",
        "fdg_adni",
        "av45",
        "abeta_num",
        "tau_num",
        "ptau_num",
        "hippocampus",
        "wholebrain",
    ]:
        if col not in enriched.columns:
            continue
        non_missing = int(enriched[col].notna().sum())
        rows.append({"field": col, "non_missing": non_missing, "coverage": non_missing / total})
    report = pd.DataFrame(rows)
    report.to_csv(out_csv, index=False)

    label_counts = {}
    for col in ["clinical_label", "dx", "dx_bl", "conversion_label"]:
        if col in enriched.columns:
            label_counts[col] = enriched[col].dropna().astype(str).value_counts().to_dict()

    lines = [
        "# ADNI Metadata Match Report",
        "",
        f"- PET/SUVR samples: {total}",
        f"- Matched clinical rows: {int(enriched['clinical_label'].notna().sum())}",
        "",
        "## Field Coverage",
        "",
        "| field | non_missing | coverage |",
        "|---|---:|---:|",
    ]
    for row in rows:
        lines.append(f"| {row['field']} | {row['non_missing']} | {row['coverage']:.3f} |")
    lines.extend([
        "",
        "## Label Counts",
        "",
    ])
    for col, counts in label_counts.items():
        lines.append(f"### {col}")
        lines.append("")
        for key, value in counts.items():
            lines.append(f"- {key}: {value}")
        lines.append("")
    out_md.write_text("\n".join(lines), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Match ADNI Excel metadata to the PET/SUVR manifest.")
    parser.add_argument("--manifest", type=Path, default=Path("data/metadata/splits/pet_fdg_manifest_with_split.csv"))
    parser.add_argument("--full-xlsx", type=Path, default=Path("data/ADNIbase1416_info.xlsx"))
    parser.add_argument("--conversion-xlsx", type=Path, default=Path("data/adni_1203s_info_fix.xlsx"))
    parser.add_argument("--out", type=Path, default=Path("data/metadata/adni_matched_clinical.csv"))
    parser.add_argument("--split-out-dir", type=Path, default=Path("data/metadata/splits"))
    parser.add_argument("--report-md", type=Path, default=Path("data/metadata/adni_match_report.md"))
    parser.add_argument("--report-csv", type=Path, default=Path("data/metadata/adni_match_report.csv"))
    args = parser.parse_args()

    manifest = pd.read_csv(args.manifest)
    manifest["subject_id_norm"] = manifest["subject_id"].map(normalize_subject_id)

    full = read_xlsx(args.full_xlsx, FULL_COLUMNS)
    full = add_numeric_clean_columns(full)
    full = full.rename(
        columns={
            "label": "clinical_label",
            "DX": "dx",
            "DX_bl": "dx_bl",
            "AGE": "age",
            "PTGENDER": "sex",
            "PTEDUCAT": "education",
            "PTETHCAT": "ethnicity",
            "PTRACCAT": "race",
            "PTMARRY": "marital_status",
            "APOE4": "apoe4",
            "FDG": "fdg_adni",
            "PIB": "pib",
            "AV45": "av45",
            "ABETA": "abeta",
            "TAU": "tau",
            "PTAU": "ptau",
            "ABETA_num": "abeta_num",
            "TAU_num": "tau_num",
            "PTAU_num": "ptau_num",
            "CDRSB": "cdrsb",
            "ADAS11": "adas11",
            "ADAS13": "adas13",
            "ADASQ4": "adasq4",
            "MMSE": "mmse",
            "RAVLT_immediate": "ravlt_immediate",
            "RAVLT_learning": "ravlt_learning",
            "RAVLT_forgetting": "ravlt_forgetting",
            "RAVLT_perc_forgetting": "ravlt_perc_forgetting",
            "LDELTOTAL": "ldeltotal",
            "DIGITSCOR": "digitscor",
            "TRABSCOR": "trabscor",
            "FAQ": "faq",
            "MOCA": "moca",
            "Years_bl": "years_bl",
        }
    )

    conversion = read_xlsx(args.conversion_xlsx, CONVERSION_COLUMNS)
    conversion = conversion.rename(
        columns={
            "label": "conversion_label",
            "DX.1": "dx_followup",
            "mPACCdigit": "mpacc_digit",
            "mPACCtrailsB": "mpacc_trailsb",
            "Ventricles（心室）": "ventricles",
            "Hippocampus（海马）": "hippocampus",
            "WholeBrain（全脑）": "wholebrain",
            "Entorhinal（内嗅觉）": "entorhinal",
            "Fusiform（梭形）": "fusiform",
            "MidTemp（中点温度）": "midtemp",
        }
    )
    conversion = conversion.drop(columns=["PTID"], errors="ignore")

    enriched = manifest.merge(full.drop(columns=["PTID"], errors="ignore"), on="subject_id_norm", how="left")
    enriched = enriched.merge(conversion, on="subject_id_norm", how="left")
    enriched = enriched.drop(columns=["subject_id_norm"])

    args.out.parent.mkdir(parents=True, exist_ok=True)
    enriched.to_csv(args.out, index=False)
    write_split_manifests(enriched, args.split_out_dir)
    summarize(enriched, args.report_md, args.report_csv)

    print(f"wrote={args.out}")
    print(f"wrote_splits={args.split_out_dir}/*_clinical.csv")
    print(f"wrote_report={args.report_md}")
    print(f"samples={len(enriched)} matched={int(enriched['clinical_label'].notna().sum())}")


if __name__ == "__main__":
    main()