"""Run code/preprocess.py on every subject in data/ds002748, in parallel processes.

Usage: python code/preprocess_all.py <out dir>
"""
import glob
import os
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor

ATLAS = "data/atlas/Schaefer2018_100Parcels_7Networks_order_FSLMNI152_2mm.nii.gz"


def one(bold, out):
    sub = os.path.basename(bold).split("_")[0]
    prefix = os.path.join(out, sub)
    if os.path.exists(prefix + "_qc.json"):
        return sub, 0
    env = dict(os.environ, OMP_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1", MKL_NUM_THREADS="1",
               PYTHONWARNINGS="ignore")
    r = subprocess.run([sys.executable, "code/preprocess.py", bold, ATLAS, prefix], env=env,
                       stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    if r.returncode:
        print(r.stdout[-2000:], file=sys.stderr)
    return sub, r.returncode


def main():
    out = sys.argv[1]
    os.makedirs(out, exist_ok=True)
    bolds = sorted(glob.glob("data/ds002748/sub-*/func/sub-*_task-rest_bold.nii.gz"))
    workers = max(1, min(len(bolds), (os.cpu_count() or 1)))
    with ThreadPoolExecutor(workers) as ex:
        for sub, rc in ex.map(lambda b: one(b, out), bolds):
            print(sub, "ok" if rc == 0 else f"failed ({rc})", flush=True)


if __name__ == "__main__":
    main()