import os
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

DATA_DIR = "/home/wenting/zr/Segmentation/data/lidc_TotalSeg"
EXPECTED_COUNT = 3000

def check_patient(patient_folder):
    patient_dir = os.path.join(DATA_DIR, patient_folder)
    if not os.path.isdir(patient_dir):
        return None
        
    img_dir = os.path.join(patient_dir, "04_drr_256", "cxr")
    mask_dir = os.path.join(patient_dir, "04_drr_256", "mask_compact")
    meta_path = os.path.join(patient_dir, "02_totalseg", "phase2_metadata.json")
    
    status = {"patient": patient_folder, "imgs": 0, "masks": 0, "has_meta": False, "issue": ""}
    
    # 1. 检查 Metadata
    if os.path.exists(meta_path):
        status["has_meta"] = True
    else:
        status["issue"] += "Missing metadata JSON. "
        
    # 2. 检查图片数量
    if os.path.exists(img_dir):
        # 快速统计文件数
        status["imgs"] = len([f for f in os.scandir(img_dir) if f.name.endswith('.png')])
    else:
        status["issue"] += "Missing cxr dir. "
        
    # 3. 检查Mask数量
    if os.path.exists(mask_dir):
        status["masks"] = len([f for f in os.scandir(mask_dir) if f.name.endswith('.npz')])
    else:
        status["issue"] += "Missing mask dir. "
        
    if status["imgs"] != EXPECTED_COUNT or status["masks"] != EXPECTED_COUNT or not status["has_meta"]:
        return status
    return None

def main():
    patients = sorted([d for d in os.listdir(DATA_DIR) if d.startswith("LIDC-IDRI-")])
    print(f"[*] Scanning {len(patients)} patients in {DATA_DIR}...")
    
    anomalies =[]
    # 使用多线程加速扫描
    with ThreadPoolExecutor(max_workers=16) as executor:
        results = list(tqdm(executor.map(check_patient, patients), total=len(patients)))
        
    for res in results:
        if res is not None:
            anomalies.append(res)
            
    print("\n" + "="*50)
    print(f"[*] Found {len(anomalies)} patients with incomplete data:")
    print("="*50)
    
    for a in anomalies:
        issue_str = a["issue"] if a["issue"] else f"Incomplete pairs (Imgs: {a['imgs']}, Masks: {a['masks']})"
        print(f"- {a['patient']}: {issue_str}")

if __name__ == "__main__":
    main()