"""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"))
    # each process peaks near 0.7 GB; leave headroom so no subject is killed for memory
    try:
        mem_gb = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / 1e9
    except (ValueError, OSError):
        mem_gb = 4.0
    try:
        with open("/sys/fs/cgroup/memory.max") as fh:
            v = fh.read().strip()
            if v != "max":
                mem_gb = min(mem_gb, int(v) / 1e9)
    except OSError:
        pass
    workers = max(1, min(len(bolds), os.cpu_count() or 1, int(mem_gb // 1.0)))
    print("workers", workers, flush=True)
    failed = []
    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 rc:
                failed.append(sub)
    # retry failures one at a time
    for b in bolds:
        sub = os.path.basename(b).split("_")[0]
        if sub in failed:
            sub, rc = one(b, out)
            print(sub, "retry ok" if rc == 0 else f"retry failed ({rc})", flush=True)
            if rc:
                sys.exit(1)


if __name__ == "__main__":
    main()