dataLoaderLazyADNI.py

import json
from pathlib import Path

import pandas as pd
import os
PROJECT_DIR=os.environ.get("PROJECT_DIR", "/scratch/progetto")

def load_subjects_from_jsonl(jsonl_path):
    subjects = []
    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line in f:
            subjects.append(json.loads(line))
    return subjects
def load_adni(forModelTraining=True):
    if forModelTraining:
        train = load_subjects_from_jsonl(PROJECT_DIR+"/train_paths_subjects_adni.jsonl")
    else:
        train=load_subjects_from_jsonl(PROJECT_DIR+"/test_paths_subjects_adni.jsonl")
    print(f"✅ Loaded {len(train)} training subjects.")
    return train

def dataLoaderLazyADNI(
    root_dir,
    output_file,
    meta_path,
    pattern="*preproc-quasiraw*.nii.gz",
    internal_col="Image Data ID",     # <- internal ID (es. I384810)
    subject_col="Subject"      # <- subject ID (es. 018_S_5250)
):
    print(">>> Inizio dataLoader per ADNI")
    print("Working dir:", Path.cwd())
    print("Scanning root:", root_dir)

    # Load metadata
    try:
        meta_df = pd.read_csv(meta_path, dtype=str)
        print(f"✅ Loaded metadata with {len(meta_df)} rows.")
    except Exception as e:
        print(f"❌ Errore caricamento file metadata: {e}")
        return

    if internal_col not in meta_df.columns or subject_col not in meta_df.columns:
        print(f"❌ Il CSV deve contenere le colonne '{internal_col}' e '{subject_col}'")
        return

    # Scan for .nii.gz files (exclude mask)
    nii_files = sorted([f for f in root_dir.rglob(pattern) if "mask" not in f.name])
    print(f"🔍 Found {len(nii_files)} matching .nii.gz files (escludendo mask).")

    if not nii_files:
        print("⚠️ Nessun file trovato! Verifica il path e il pattern.")
        return

    entries = []
    direct_matches, dir_matches, no_matches = 0, 0, 0
    scratch_base = Path.home()
    for nii_path in nii_files:
        internal_id = nii_path.name.split("_")[0]  # es: "I384810"
        rel_path = "/"+str(nii_path.relative_to(scratch_base))
        entry = {
            "internal_id": internal_id,
            "image_path": rel_path
        }

        # Primo tentativo: match diretto su Image Data ID
        row = meta_df.loc[meta_df[internal_col] == internal_id]
        if not row.empty:

            entry["sex"] = row.iloc[0]["Sex"]
            entry["age"] = row.iloc[0]["Age"]
            entry["group"] = row.iloc[0]["Group"]
            direct_matches += 1

        else:
            # Secondo tentativo: risalire al participant dalle cartelle nella stessa dir
            parent_dir = nii_path.parent
            matched_subject = None
            for d in parent_dir.iterdir():
                if d.is_dir() and internal_id in d.name:
                    matched_subject = "_".join(d.name.split("_")[:3])  # es: "018_S_5250"
                    break

            if matched_subject:
                row = meta_df.loc[meta_df[subject_col] == matched_subject]
                if not row.empty:

                    entry["sex"] = row.iloc[0]["Sex"]
                    entry["age"] = row.iloc[0]["Age"]
                    entry["group"] = row.iloc[0]["Group"]
                    dir_matches += 1
                    print(f"🔄 Mappato {internal_id} -> {matched_subject} tramite cartella")
                else:
                    print(f"⚠️ Nessuna metadata trovata per participant_id {matched_subject}")

                    entry["sex"] = None
                    entry["age"] = None
                    entry["group"] = None
                    no_matches += 1
            else:
                print(f"⚠️ Metadata non trovata per internal_id {internal_id}")

                entry["sex"] = None
                entry["age"] = None
                entry["group"] = None
                no_matches += 1

        entries.append(entry)

    print(f"\n📊 Report finale:")
    print(f"  ✅ Match diretti (Image Data ID): {direct_matches}")
    print(f"  🔄 Match tramite cartella (participant_id): {dir_matches}")
    print(f"  ⚠️ Non trovati: {no_matches}")

    print(f"\n✅ Writing {len(entries)} entries to: {output_file}")
    with open(output_file, "w", encoding="utf-8") as f:
        for entry in entries:
            json.dump(entry, f)
            f.write("\n")

    return None


if __name__ == '__main__':
    dataLoaderLazyADNI(
        root_dir=Path.home() / "scratch" / "ADNI",
        output_file="paths_subjects_adni.jsonl",
        meta_path=Path("adni_labels.csv"),
        internal_col="Image Data ID",
        subject_col="Subject"
    )