pesky-fieldmaps/new_renamer.py

from enum import Enum
from itertools import product
import json
import logging
from pathlib import Path
import re
from typing import NamedTuple

from bids import BIDSLayout
import click
import numpy as np
import numpy.typing as npt
from rich.highlighter import NullHighlighter
from rich.logging import RichHandler

# logging will help us see where things stand out from the norm
FORMAT = "%(message)s"
logging.basicConfig(
    level=logging.INFO,
    format=FORMAT,
    datefmt="[%X]",
    handlers=[RichHandler(highlighter=NullHighlighter(), keywords=[])],
)
log = logging.getLogger("rich")

# when setting IntendedFor we may want to distinguish several states
# to inform mitigation (and logging)
IntendedForStatus = Enum(
    "IntendedForStatus",
    [
        ("NOTHING_TO_DO", 0),
        ("CHANGES_REQUIRED", 1),
        ("DONE", 2),
        ("SOURCE_AND_TARGETS_DO_NOT_EXIST", 3),
        ("SOURCE_DOES_NOT_EXIST", 4),
        ("TARGET_DOES_NOT_EXIST", 5),
        ("SHIMS_DONT_MATCH", 6),
        ("FILE_NOT_READABLE", 7),
        ("FILE_NOT_WRITABLE", 8),
    ],
)

# building paths with f-strings is hard, so let's store dicts of BIDS
# entities to define renamings, deletions, and IntendedFor; named
# tuples (typed version) will be used to make things more explicit

type dbent = dict[str, str | int | None]  # a dict of bids entities, {name: value}


class RenameEntitySpec(NamedTuple):
    match: dbent
    replace: dbent
    void_ok: bool = False


class DeleteEntitySpec(NamedTuple):
    match: dbent


class IntendedForSpec(NamedTuple):
    fmap: dbent
    target: list[dbent]


# for renames and deletes, we will concretize these as paths
class RenamePathSpec(NamedTuple):
    before: Path
    after: Path


class DeletePathSpec(NamedTuple):
    path: Path


@click.command()
@click.argument(
    "bids_root",
    type=click.Path(
        exists=True, file_okay=False, dir_okay=True, resolve_path=True, path_type=Path
    ),
)
@click.argument("subjects", nargs=-1)
@click.option("--session", default="001", show_default=True)
@click.option(
    "--pause", is_flag=True, help="Pause after each subject (to take in logs)"
)
@click.option(
    "--cache",
    type=click.Path(file_okay=False, dir_okay=True, resolve_path=True, path_type=Path),
)
def main(
    bids_root: Path,
    subjects: tuple[str, ...],
    session: str,
    cache: Path | None,
    pause: bool,
):
    ds = BIDSLayout(bids_root, validate=False, database_path=cache)

    if len(subjects) == 0:
        subjects = tuple(ds.get_subjects())

    for subject in subjects:
        log.info("Processing subject %s", subject)
        apply_renames(ops, ds, subject, session)
        reassign_func_fieldmaps(func_intentions, ds, subject, session)
        reassign_dwi_fieldmaps(dwi_intentions, ds, subject, session)

        if pause:
            input("Hit enter to proceed to next subject")


def save_pretty_json(filename: Path, data: dict) -> None:
    j = json_dumps_pretty(data)
    with filename.open("w") as fp:
        fp.write(j)


def json_dumps_pretty(j: dict, indent: int = 2, sort_keys: bool = True) -> str:
    """Given a json structure, pretty print it by colliding numeric arrays
    into a line.

    If resultant structure differs from original -- throws exception.

    Copied & minimally changed from heudiconv 1.4.0 (utils.py).
    https://github.com/nipy/heudiconv/

    Copyright HeuDiConv developers; licensed under the Apache License,
    Version 2.0. http://www.apache.org/licenses/LICENSE-2.0

    """
    js = json.dumps(j, indent=indent, sort_keys=sort_keys)
    # trim away \n and spaces between entries of numbers
    js_ = re.sub(
        '[\n ]+("?[-+.0-9e]+"?,?) *\n(?= *"?[-+.0-9e]+"?)',
        r" \1",
        js,
        flags=re.MULTILINE,
    )
    # uniform no spaces before ]
    js_ = re.sub(r" *\]", "]", js_)
    # uniform spacing before numbers
    # But that thing could screw up dates within strings which would have 2 spaces
    # in a date like Mar  3 2017, so we do negative lookahead to avoid changing
    # in those cases
    js_ = re.sub(
        r"(?<!\w{3})"  # negative lookbehind for the month
        r'  *("?[-+.0-9e]+"?)'
        r"(?! [123]\d{3})"  # negative lookahead for a year
        r"(?P<space> ?)[ \n]*",
        r" \1\g<space>",
        js_,
    )
    # no spaces after [
    js_ = re.sub(r"\[ ", "[", js_)
    # the load from the original dump and reload from tuned up
    # version should result in identical values since no value
    # must be changed, just formatting.
    j_just_reloaded = json.loads(js)
    j_tuned = json.loads(js_)

    assert j_just_reloaded == j_tuned, (
        "Values differed when they should have not. Report to the heudiconv developers"
    )

    return js_


def build_paths(
    layout: BIDSLayout, spec: RenameEntitySpec | DeleteEntitySpec, const: dbent
) -> RenamePathSpec | DeletePathSpec:
    """A helper to build BIDS paths from dicts

    This uses dict.update ("|" operator) to replace entities, and
    builds paths from there. The const part should contain subject and
    session.

    """
    if isinstance(spec, RenameEntitySpec):
        return RenamePathSpec(
            before=Path(layout.build_path(const | spec.match)),
            after=Path(layout.build_path(const | spec.match | spec.replace)),
        )
    if isinstance(spec, DeleteEntitySpec):
        try:
            path_to_delete = Path(layout.build_path(const | spec.match))
        except ValueError as err:
            if (ext := spec.match.get("extension")) in (".bval", ".bvec"):
                # _sbref.[bval|bvec] is so wrong that pybids refuses
                # to build it, so we ask pybids to build .json and
                # swap back the extension
                tmp = layout.build_path(const | spec.match | {"extension": ".json"})
                path_to_delete = Path(tmp).with_suffix(ext)
            else:
                raise err

        return DeletePathSpec(path=path_to_delete)


def edit_scans_file(scans_tsv: Path, repls: dict[Path, Path]) -> None:
    try:
        txt = scans_tsv.read_text()
    except FileNotFoundError:
        log.error("%s file is missing, did you get it?", scans_tsv)
        return

    for old, new in repls.items():
        txt = txt.replace(str(old), str(new))

    try:
        scans_tsv.write_text(txt)
    except PermissionError:
        log.error("%s permission error, did you unlock it?", scans_tsv)


def apply_renames(
    specs: list[RenameEntitySpec | DeleteEntitySpec],
    layout: BIDSLayout,
    subject: str,
    session: str,
) -> None:
    def rel2ss(p):
        return p.relative_to(ss_dir)

    ss_dir = Path(layout.root) / f"sub-{subject}" / f"ses-{session}"
    scans_file_path = ss_dir / f"sub-{subject}_ses-{session}_scans.tsv"
    repls = {}

    for s in specs:
        spec = build_paths(layout, s, {"subject": subject, "session": session})

        if isinstance(spec, DeletePathSpec):  # if spec.after is None
            try:
                spec.path.unlink()
                log.info("removed %s", spec.path.relative_to(layout.root))
            except FileNotFoundError:
                log.info(
                    "nothing to do, %s does not exist",
                    spec.path.relative_to(layout.root),
                )
            continue

        source_exists = spec.before.exists() or spec.before.is_symlink()
        target_exists = spec.after.exists() or spec.after.is_symlink()

        if target_exists and not source_exists:
            log.info("nothing to do; %s already exists", spec.after.name)
            continue  # perfect, no-op
        elif not (source_exists or target_exists):
            # a "void" operation
            if isinstance(s, RenameEntitySpec) and s.void_ok:
                # this was marked as expected
                log.info(
                    "void op; requested: %s → %s; neither exists",
                    rel2ss(spec.before),
                    rel2ss(spec.after),
                )
            else:
                # we'd expect one to exist
                log.warning(
                    "requested: %s → %s; neither exists",
                    rel2ss(spec.before),
                    rel2ss(spec.after),
                )
            continue  # no-op, expected or not
        elif target_exists:
            log.warning("clobbering %s", rel2ss(spec.after))
            # not ideal, we operate nevertheless

        log.info("%s → %s", rel2ss(spec.before), rel2ss(spec.after))

        _ = spec.before.rename(spec.after)
        if spec.before.suffix == ".gz" or spec.before.suffix == ".nii":
            repls[spec.before.relative_to(ss_dir)] = spec.after.relative_to(ss_dir)

    if len(repls) > 0:
        edit_scans_file(scans_file_path, repls)


def load_shim(p: Path) -> npt.NDArray:
    with p.open() as fp:
        d = json.load(fp)
    return np.array(d["ShimSetting"])


def shims_match(paths: list[Path]) -> np.bool:
    shims = np.array([load_shim(p) for p in paths])
    return np.all(shims == shims[0])


def ensure_intended_for(
    layout: BIDSLayout,
    fmap_entities: dbent,
    target_entities: list[dbent],
    subject: str,
    session: str,
    dry_run: bool = False,
) -> IntendedForStatus:
    const = {"subject": subject, "session": session}
    fmap_sidecar = Path(
        layout.build_path(fmap_entities | const | {"extension": ".json"})
    )
    target_sidecars = [
        Path(layout.build_path(entities | const | {"extension": ".json"}))
        for entities in target_entities
    ]

    # targets: forward-slash separated paths relative to the
    # participant subdirectory are deprecated (as of BIDS 1.11.1) but
    # I'm not sure about the support for bids:: notation in qsiprep;
    # heudiconv uses the relative path anyway
    targets = [
        Path(layout.build_path(entities | const))
        .relative_to(Path(layout.root) / f"sub-{subject}")
        .as_posix()
        for entities in target_entities
    ]

    # This function gets a statement: this should be intended for
    # that. There may be nothing to do - either because things are
    # already in a desired state, or because we can not enact the
    # setting. We'll go through checks in order.

    # we are given paths which may or may not exist
    if not (fmap_sidecar.exists() or fmap_sidecar.is_symlink()):
        # source does not exist, but maybe targets don't exist either
        for target in target_sidecars:
            if target.exists() or target.is_symlink():
                # nope, at least one target exists
                return IntendedForStatus.SOURCE_DOES_NOT_EXIST
        # neither source nor targets exists, this is less of an issue
        return IntendedForStatus.SOURCE_AND_TARGETS_DO_NOT_EXIST

    for target in target_sidecars:
        # here we are could be interested in two things: whether the
        # target acquisition exists at all (using sidecar as a proxy
        # for the .nii.gz file - not 100% correct but good enough for
        # our case) and whether the file is readable; for now we will
        # only check existence...
        if not (target.exists() or target.is_symlink()):
            return IntendedForStatus.TARGET_DOES_NOT_EXIST

    # all shims have to match
    try:
        shims_ok = shims_match([fmap_sidecar, *target_sidecars])
    except FileNotFoundError:
        return IntendedForStatus.FILE_NOT_READABLE
    if not shims_ok:
        return IntendedForStatus.SHIMS_DONT_MATCH

    # read the field map sidecar to check its current IntendedFor
    try:
        with fmap_sidecar.open() as jp:
            j = json.load(jp)
    except FileNotFoundError:
        return IntendedForStatus.FILE_NOT_READABLE

    # perhaps all targets are OK and this is a no-op
    if j.get("IntendedFor") == targets:
        return IntendedForStatus.NOTHING_TO_DO
    elif dry_run:
        return IntendedForStatus.CHANGES_REQUIRED
    else:
        j["IntendedFor"] = targets

    # if changes are required and it's not a dry run, apply them
    try:
        save_pretty_json(fmap_sidecar, j)
    except PermissionError:
        return IntendedForStatus.FILE_NOT_WRITABLE

    return IntendedForStatus.DONE


def reassign_func_fieldmaps(
    func_intentions: list[IntendedForSpec],
    layout: BIDSLayout,
    subject: str,
    session: str,
) -> None:
    # Start with dry run for task field maps (because initial
    # assignment might be two-to-one, we don't want to overwrite
    # unless we know we can unambiguously assign others)

    results = []
    for intention in func_intentions:
        results.append(
            ensure_intended_for(
                layout=layout,
                fmap_entities=intention.fmap,
                target_entities=intention.target,
                subject=subject,
                session=session,
                dry_run=True,
            )
        )

    if all(res == IntendedForStatus.SOURCE_DOES_NOT_EXIST for res in results):
        log.warning("sub-%s: no func-related field maps to reassign", subject)
    elif all(res == IntendedForStatus.NOTHING_TO_DO for res in results):
        log.info("sub-%s: all func-related field maps assigned correctly", subject)
    elif any(res.value > IntendedForStatus.DONE.value for res in results):
        issues_found = set(res.name for res in results)
        log.warning(
            "sub-%s: func-related field map preset can't be applied: %s",
            subject,
            issues_found,
        )
    else:
        for intention in func_intentions:
            res = ensure_intended_for(
                layout=layout,
                fmap_entities=intention.fmap,
                target_entities=intention.target,
                subject=subject,
                session=session,
                dry_run=False,
            )
            if res.value > IntendedForStatus.DONE.value:
                # we got here after dry run, likely r/w issue, present as error
                log.error(
                    "Failed when assigning field map %s, %s", intention.fmap, res.name
                )
            else:
                log.info("Assigning field map %s, %s", intention.fmap, res.name)


def reassign_dwi_fieldmaps(
    dwi_intentions: list[IntendedForSpec],
    layout: BIDSLayout,
    subject: str,
    session: str,
) -> None:
    # for dwi, the files are independent and we are OK if an
    # acquisition is missing, so no dry runs
    for intention in dwi_intentions:
        res = ensure_intended_for(
            layout=layout,
            fmap_entities=intention.fmap,
            target_entities=intention.target,
            subject=subject,
            session=session,
        )
        if res == IntendedForStatus.SOURCE_AND_TARGETS_DO_NOT_EXIST:
            log.info("nothing to do, %s, %s", intention.fmap, res.name)
        elif res.value > IntendedForStatus.DONE.value:
            log.error(
                "Failed when assigning field map %s, %s", intention.fmap, res.name
            )
        else:
            log.info("Assigning field map %s, %s", intention.fmap, res.name)


## Rename & delete specifications

ops = []

# the original run-1/2/3 become acq-rest/restvidn/restvida
label_for_run = {1: "rest", 2: "restvidn", 3: "restvida"}
for dir, run, ext in product(("ap", "pa"), (1, 2, 3), (".nii.gz", ".json")):
    ops.append(
        RenameEntitySpec(
            match={
                "datatype": "fmap",
                "direction": dir,
                "run": run,
                "suffix": "epi",
                "extension": ext,
            },
            replace={"run": None, "acquisition": label_for_run[run]},
        )
    )

# dwi files without dir-<label> should have been dir-pa
for acq, ext in product(("b1200", "mshell"), (".nii.gz", ".json")):
    ops.append(
        RenameEntitySpec(
            match={
                "datatype": "dwi",
                "acquisition": acq,
                "direction": None,
                "suffix": "dwi",
                "extension": ext,
            },
            replace={"direction": "PA"},
            void_ok=(acq == "mshell"),
        )
    )

# dwi files with sbref suffix should have gone into fmap
for acq, ext in product(("b1200", "mshell"), (".nii.gz", ".json")):
    ops.append(
        RenameEntitySpec(
            match={
                "datatype": "dwi",
                "acquisition": acq,
                "suffix": "sbref",
                "extension": ext,
            },
            replace={"datatype": "fmap", "direction": "AP", "suffix": "epi"},
            void_ok=(acq == "mshell"),
        )
    )

# ... and their bval / bvec should have never existed
for acq, ext in product(("b1200", "mshell"), (".bval", ".bvec")):
    ops.append(
        DeleteEntitySpec(
            match={
                "datatype": "dwi",
                "acquisition": acq,
                "suffix": "sbref",
                "extension": ext,
            },
        )
    )


# IntendedFor additions and corrections
func_intentions = []  # we'll want to apply these as a set
dwi_intentions = []  # we'll want to apply these independently

for dir, task in product(("ap", "pa"), ("rest", "restvidn", "restvida")):
    func_intentions.append(
        IntendedForSpec(
            fmap={
                "datatype": "fmap",
                "direction": dir,
                "acquisition": task,
                "suffix": "epi",
            },
            target=[{"datatype": "func", "task": task, "suffix": "bold"}],
        )
    )

for acq in ("b1200", "mshell"):
    dwi_intentions.append(
        IntendedForSpec(
            fmap={
                "datatype": "fmap",
                "direction": "AP",
                "acquisition": acq,
                "suffix": "epi",
            },
            target=[
                {
                    "datatype": "dwi",
                    "direction": "PA",
                    "acquisition": acq,
                    "suffix": "dwi",
                }
            ],
        )
    )


if __name__ == "__main__":
    main()