data/export_split_by_pathology.py

"""
export_split_by_pathology.py
----------------------------
Tải 1 split (mặc định: test) của MIMIC-CXR_resized từ HF, giải nén tar shards,
rồi đổ ảnh vào thư mục theo 14 nhãn bệnh lý CheXpert.

Vì sao phải làm thế này:
  - Trên HF, MIMIC-CXR_resized lưu ảnh dưới dạng tar shards (cxr-0000.tar, ...),
    KHÔNG tách theo split. Ảnh của 1 split nằm rải khắp các shard.
  - manifest_{train,val,test}.csv mới là nơi biết ảnh nào thuộc split nào, và
    chứa 14 cột `chex_<Pathology>` (giá trị U-MultiClass: 1/0/-1/blank).
  - Mỗi study là MULTI-LABEL: 1 ảnh có thể dương tính nhiều bệnh -> được copy
    vào NHIỀU thư mục (mỗi nhãn positive 1 thư mục).

Quy ước nhãn (giống mimic_cxr_resized_builder._row_to_pnu):
    "1" / "1.0"   -> positive  -> bỏ vào thư mục <Pathology>/
    "-1" / "-1.0" -> uncertain -> tuỳ --uncertain
    "0"/"0.0"/blank -> negative -> bỏ qua

Cấu trúc output (mặc định --uncertain separate):
    OUT/
      Cardiomegaly/<dicom>.jpg
      Pleural_Effusion/<dicom>.jpg
      No_Finding/<dicom>.jpg
      _uncertain/Atelectasis/<dicom>.jpg     (nếu --uncertain separate)
      _summary.csv                            (đếm ảnh mỗi nhãn)

Chạy LOCAL (đã có data) hoặc trên Colab/Kaggle (tự tải từ HF):
    # tải từ HF rồi export test
    python data/export_split_by_pathology.py --out ./test_by_pathology

    # đã giải nén sẵn cây files/ ở đâu đó -> bỏ qua tải/giải nén
    python data/export_split_by_pathology.py --out ./test_by_pathology \
        --extracted_root /content/data/MIMIC-CXR_resized \
        --manifest /content/data/MIMIC-CXR_resized/manifest_test.csv

    # export cả split val, dùng hardlink cho đỡ tốn ổ
    python data/export_split_by_pathology.py --split val --link hardlink
"""
from __future__ import annotations

import argparse
import csv
import os
import sys
import tarfile
from collections import defaultdict
from pathlib import Path

# 14 nhãn CheXpert — single source of truth.
try:
    from model.chexpert_classifier import PATHOLOGIES
except Exception:
    # fallback nếu chạy ngoài project (giữ đúng thứ tự).
    PATHOLOGIES = [
        "No Finding", "Enlarged Cardiomediastinum", "Cardiomegaly", "Lung Opacity",
        "Lung Lesion", "Edema", "Consolidation", "Pneumonia", "Atelectasis",
        "Pneumothorax", "Pleural Effusion", "Pleural Other", "Fracture",
        "Support Devices",
    ]

_POS = {"1", "1.0"}
_UNC = {"-1", "-1.0"}


def _safe(name: str) -> str:
    """Tên thư mục an toàn: 'Pleural Effusion' -> 'Pleural_Effusion'."""
    return name.replace(" ", "_")


def _norm(p: str) -> str:
    """Chuẩn hoá path để so khớp tar member name <-> manifest image_relpath."""
    return p.replace("\\", "/").lstrip("/")


# ── Phase 1: tải từ HF (manifest + shards) ──────────────────────────────────

def download_from_hf(repo_id: str, split: str, work: Path) -> tuple[Path, list[Path]]:
    """Tải manifest_<split>.csv + toàn bộ tar shards về `work`.
    Trả về (manifest_path, [shard_paths])."""
    from huggingface_hub import snapshot_download

    manifest_name = {"train": "manifest_train.csv",
                     "val": "manifest_val.csv",
                     "validate": "manifest_val.csv",
                     "test": "manifest_test.csv"}[split]

    print(f"[download] snapshot_download {repo_id}:MIMIC-CXR_resized "
          f"(manifest + shards) -> {work}")
    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        local_dir=str(work),
        allow_patterns=[
            f"MIMIC-CXR_resized/{manifest_name}",
            "MIMIC-CXR_resized/shards/*.tar",
        ],
    )
    mr = work / "MIMIC-CXR_resized"
    manifest = mr / manifest_name
    shards = sorted((mr / "shards").glob("*.tar"))
    if not manifest.is_file():
        sys.exit(f"ERROR: không thấy manifest sau khi tải: {manifest}")
    if not shards:
        sys.exit(f"ERROR: không thấy tar shard nào dưới {mr/'shards'}")
    print(f"[download] manifest={manifest.name}  shards={len(shards)}")
    return manifest, shards


# ── Phase 2: đọc manifest -> map ảnh test -> nhãn ───────────────────────────

def load_label_map(manifest: Path):
    """Trả về dict: image_relpath(norm) -> {'pos': set, 'unc': set, 'report': str|None}."""
    label_map: dict[str, dict] = {}
    missing_cols = None
    with open(manifest, encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        cols = reader.fieldnames or []
        chex_cols = {p: f"chex_{p}" for p in PATHOLOGIES if f"chex_{p}" in cols}
        missing_cols = [p for p in PATHOLOGIES if f"chex_{p}" not in cols]
        rel_col = "image_relpath" if "image_relpath" in cols else None
        if rel_col is None:
            sys.exit(f"ERROR: manifest thiếu cột 'image_relpath'. Có: {cols}")
        has_report = "report_relpath" in cols
        for row in reader:
            rel = _norm(str(row[rel_col]).strip())
            pos, unc = set(), set()
            for path, col in chex_cols.items():
                v = str(row.get(col, "")).strip()
                if v in _POS:
                    pos.add(path)
                elif v in _UNC:
                    unc.add(path)
            rep = _norm(str(row["report_relpath"]).strip()) if has_report else None
            label_map[rel] = {"pos": pos, "unc": unc, "report": rep or None}
    if missing_cols:
        print(f"[labels] CẢNH BÁO: manifest thiếu cột cho: {missing_cols}")
    if not has_report:
        print("[labels] CẢNH BÁO: manifest không có cột 'report_relpath' → bỏ qua report")
    print(f"[labels] {len(label_map):,} ảnh trong manifest")
    return label_map


def gather_reports(shards: list[Path], report_set: set) -> dict:
    """Pass phụ: rút text của các report cần dùng từ tar (report nằm rải, gom 1 lượt).
    Report là file .txt nhỏ nên giữ trong RAM thoải mái."""
    reports: dict[str, bytes] = {}
    if not report_set:
        return reports
    for shard in shards:
        with tarfile.open(shard, "r") as tf:
            for m in tf:
                if not m.isfile():
                    continue
                name = _norm(m.name)
                if name in report_set and name not in reports:
                    reports[name] = tf.extractfile(m).read()
    print(f"[reports] rút được {len(reports):,} / {len(report_set):,} report")
    return reports


# ── Phase 3: rút ảnh từ tar -> thư mục theo nhãn ────────────────────────────

def _place(data: bytes, dicom_name: str, paths: set, base: Path,
           counts: defaultdict, link_mode: str, report: bytes | None = None):
    """Ghi 1 ảnh (và report cùng tên .txt nếu có) vào nhiều thư mục nhãn."""
    txt_name = Path(dicom_name).stem + ".txt"
    first_written: Path | None = None
    for lab in paths:
        d = base / _safe(lab)
        d.mkdir(parents=True, exist_ok=True)
        dst = d / dicom_name
        counts[lab] += 1
        # report .txt đặt cạnh ảnh, cùng tên
        if report is not None:
            (d / txt_name).write_bytes(report)
        if dst.exists():
            continue
        if link_mode == "copy" or first_written is None:
            dst.write_bytes(data)
            first_written = dst
        else:
            try:
                if link_mode == "hardlink":
                    os.link(first_written, dst)
                else:  # symlink
                    os.symlink(os.path.abspath(first_written), dst)
            except OSError:
                dst.write_bytes(data)   # fallback nếu FS không hỗ trợ link


def export(shards: list[Path], label_map: dict, out: Path,
           uncertain: str, link_mode: str, with_report: bool = True):
    out.mkdir(parents=True, exist_ok=True)
    unc_base = out / "_uncertain"

    test_set = set(label_map.keys())

    # Gom report cần dùng (1 pass phụ qua tar) trước khi rút ảnh.
    reports: dict = {}
    if with_report:
        report_set = {label_map[k]["report"] for k in test_set
                      if label_map[k].get("report")}
        reports = gather_reports(shards, report_set)

    counts_pos: defaultdict = defaultdict(int)
    counts_unc: defaultdict = defaultdict(int)
    n_imgs = 0
    n_no_report = 0
    seen: set[str] = set()

    for si, shard in enumerate(shards, 1):
        print(f"[extract] [{si}/{len(shards)}] {shard.name}")
        with tarfile.open(shard, "r") as tf:
            for m in tf:
                if not m.isfile():
                    continue
                name = _norm(m.name)
                if name not in test_set:
                    continue
                seen.add(name)
                ent = label_map[name]
                pos, unc = ent["pos"], ent["unc"]
                if not pos and not (uncertain != "skip" and unc):
                    # không có nhãn positive (toàn negative) -> bỏ qua
                    if not pos:
                        continue
                data = tf.extractfile(m).read()
                dicom_name = Path(name).name
                rep = reports.get(ent.get("report")) if with_report else None
                if with_report and rep is None:
                    n_no_report += 1
                n_imgs += 1
                if pos:
                    _place(data, dicom_name, pos, out, counts_pos, link_mode, rep)
                if unc and uncertain != "skip":
                    if uncertain == "merge":
                        _place(data, dicom_name, unc, out, counts_unc, link_mode, rep)
                    else:  # separate
                        _place(data, dicom_name, unc, unc_base, counts_unc, link_mode, rep)

    if with_report and n_no_report:
        print(f"[reports] CẢNH BÁO: {n_no_report:,} ảnh không tìm thấy report → chỉ có .jpg")

    missing = test_set - seen
    print(f"\n[done] ảnh test rút được: {n_imgs:,} / {len(test_set):,} trong manifest")
    if missing:
        print(f"[done] CẢNH BÁO: {len(missing):,} ảnh trong manifest không thấy trong shard "
              f"(ví dụ: {list(missing)[:3]})")
    print(f"[done] ảnh toàn-negative (không có positive): bỏ qua")

    # _summary.csv
    summ = out / "_summary.csv"
    with open(summ, "w", encoding="utf-8", newline="") as f:
        w = csv.writer(f)
        w.writerow(["pathology", "positive_images", "uncertain_images"])
        for p in PATHOLOGIES:
            w.writerow([p, counts_pos.get(p, 0), counts_unc.get(p, 0)])
    print(f"[done] thống kê -> {summ}")
    print("\n  Nhãn                         positive   uncertain")
    for p in PATHOLOGIES:
        print(f"  {p:28s} {counts_pos.get(p,0):8d}   {counts_unc.get(p,0):8d}")


# ── CLI ─────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(
        description="Export 1 split của MIMIC-CXR_resized thành thư mục theo 14 nhãn bệnh lý.")
    ap.add_argument("--out", required=True, help="Thư mục output.")
    ap.add_argument("--split", default="test", choices=["train", "val", "validate", "test"])
    ap.add_argument("--repo_id", default="hieu3636/cxr-vlm-data")
    ap.add_argument("--work", default="./_hf_resized_dl",
                    help="Thư mục cache tải từ HF (khi không dùng --extracted_root).")
    ap.add_argument("--extracted_root", default=None,
                    help="Nếu đã có shards giải nén/tar sẵn ở local: trỏ tới thư mục "
                         "MIMIC-CXR_resized (chứa shards/*.tar). Bỏ qua bước tải HF.")
    ap.add_argument("--manifest", default=None,
                    help="Đường dẫn manifest_<split>.csv (mặc định lấy trong dữ liệu đã tải).")
    ap.add_argument("--uncertain", default="separate",
                    choices=["separate", "merge", "skip"],
                    help="Xử lý nhãn uncertain: separate=thư mục _uncertain/<P>; "
                         "merge=gộp chung <P>; skip=bỏ. Mặc định separate.")
    ap.add_argument("--link", default="copy", choices=["copy", "hardlink", "symlink"],
                    help="copy (an toàn nhất) | hardlink/symlink (tiết kiệm ổ khi 1 ảnh nhiều nhãn).")
    ap.add_argument("--no_report", action="store_true",
                    help="Không ghi report .txt cạnh ảnh (mặc định CÓ ghi).")
    a = ap.parse_args()

    out = Path(a.out)

    # 1) Lấy manifest + shards
    if a.extracted_root:
        mr = Path(a.extracted_root)
        shards = sorted((mr / "shards").glob("*.tar")) or sorted(mr.glob("*.tar"))
        if not shards:
            sys.exit(f"ERROR: không thấy *.tar dưới {mr} hoặc {mr/'shards'}")
        if a.manifest:
            manifest = Path(a.manifest)
        else:
            mname = {"train": "manifest_train.csv", "val": "manifest_val.csv",
                     "validate": "manifest_val.csv", "test": "manifest_test.csv"}[a.split]
            manifest = mr / mname
        if not manifest.is_file():
            sys.exit(f"ERROR: không thấy manifest: {manifest}")
        print(f"[local] manifest={manifest}  shards={len(shards)}")
    else:
        manifest, shards = download_from_hf(a.repo_id, a.split, Path(a.work))
        if a.manifest:
            manifest = Path(a.manifest)

    # 2) Đọc nhãn
    label_map = load_label_map(manifest)

    # 3) Rút ảnh (+ report) -> thư mục nhãn
    export(shards, label_map, out, a.uncertain, a.link, with_report=not a.no_report)
    print(f"\nXong. Output: {out.resolve()}")


if __name__ == "__main__":
    main()