"""Scan processed volumes for non-finite/corrupt data. Writes data/bad_procs.json."""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import glob
import json
import numpy as np
import SimpleITK as sitk
from concurrent.futures import ProcessPoolExecutor, as_completed


def check(f):
    try:
        x = sitk.GetArrayFromImage(sitk.ReadImage(f))
        return os.path.basename(f), bool(np.isfinite(x).all()), float(x.max()) if x.size else -1.0
    except Exception as e:  # noqa
        return os.path.basename(f), False, repr(e)


def main():
    files = sorted(glob.glob(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "proc", "*.nii.gz")))
    print("checking", len(files), flush=True)
    bad = []
    with ProcessPoolExecutor(max_workers=48) as ex:
        futs = [ex.submit(check, f) for f in files]
        for i, fu in enumerate(as_completed(futs), 1):
            k, ok, mx = fu.result()
            if not ok or not np.isfinite(mx):
                bad.append(k)
            if i % 1000 == 0:
                print(i, "checked", len(bad), "bad", flush=True)
    print("BAD volumes:", len(bad))
    for b in sorted(bad):
        print(" ", b)
    with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data", "bad_procs.json"), "w") as f:
        json.dump(sorted(bad), f, indent=1)


if __name__ == "__main__":
    main()