rdmtools/trr379_rdmtools/ni_mannheim_kjp.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "datalad",
#   "datalad-next",
#   "datalad-container",
#   "datalad-core",
#   "nibabel",
#   "pydicom",
#   "trr379-rdmtools @ git+https://hub.trr379.de/q02/rdmtools.git@refs/heads/main",
# ]
# ///

from datalad_core import repo
from datalad_core.config import ConfigItem
from datalad_core.runners import (
    CommandError,
    call_git,
    call_git_oneline,
    call_git_success,
)

from pathlib import (
    Path,
    PurePosixPath,
)

import csv
import json
import sys
import tempfile
import uuid

import nibabel as nib

import rich_click as click

from rich.progress import (
    Progress,
    track,
)

from shutil import copyfile

from trr379_rdmtools.dicom_archive import make_dicom_archive
from trr379_rdmtools.utils import (
    get_remote_dataset_state,
    has_staged_changes,
    post_failure,
    post_success,
    post_noop_result,
    shallow_clone_submodule,
)


hub_url = "https://hub.trr379.de"

dicom_org = "q01-mannheim-mrisessions"

all_dicoms_ds_url = (
    "https://hub.trr379.de/q01-mannheim/mri-dicoms"
)

bids_ds_url = (
    "https://hub.trr379.de/q01-mannheim-data/mri-raw"
)

envvar_prefix = "TRR_"


default_dataset_content = {
    ".datalad/.gitattributes":
        "config annex.largefiles=nothing\n",

    ".gitattributes":
        "* annex.backend=MD5E\n"
        "**/.git* annex.largefiles=nothing\n",
}


@click.group(
    "ni-mannheim-kjp",
    epilog=(
        "See https://hub.trr379.de/q02/rdmtools "
        "for a changelog and open issues."
    ),
)
def cli():
    """
    Connect Mannheim Q01 KJP MRI data with the TRR379 RDM system.

    This workflow follows the Mannheim adult MRI workflow but uses the
    dedicated Mannheim KJP HeuDiConv heuristic and performs KJP specific
    post conversion cleanup.

    The post conversion step:

    1. Selects the valid 4D task vid BOLD acquisition.
    2. Removes the reduced slice task vid conversion.
    3. Keeps the task vid SBRef.
    4. Adds session relative IntendedFor metadata.
    5. Does not add B0FieldIdentifier or B0FieldSource.
    """

    pass


def get_archive_dest_rpath(
    dicoms: Path,
    sub_id: str,
    ses_id: str,
):
    return f"{sub_id}-{ses_id}.tar"


def get_session_dataset_url(
    sub_id: str,
    ses_id: str,
) -> str:

    return (
        f"{hub_url}/"
        f"{dicom_org}/"
        f"{sub_id}-{ses_id}.git"
    )


def get_session_dataset_rpath(
    sub_id: str,
    ses_id: str,
) -> str:

    return (
        f"sessions/"
        f"{sub_id}/"
        f"{ses_id}"
    )


def create_dataset(
    path,
    private_annex=False,
) -> repo.Worktree:

    wt = repo.Worktree.init_at(path)

    if private_annex:

        wt.config.sources["git-local"].add(
            "annex.private",
            ConfigItem(
                "true",
                coercer=bool,
            ),
        )

    wt.init_annex()

    wt.config.sources["datalad-branch"].add(
        "datalad.dataset.id",
        ConfigItem(
            str(uuid.uuid4())
        ),
    )

    for rpath, content in default_dataset_content.items():

        (
            wt.path
            / PurePosixPath(rpath)
        ).write_text(content)

    call_git(
        ["add", "."],
        cwd=wt.path,
    )

    call_git(
        [
            "commit",
            "-q",
            "-m",
            "new datalad dataset",
        ],
        cwd=wt.path,
    )

    return wt


param_cfg = {

    "dicoms": {

        "type": click.Path(
            exists=True,
            file_okay=False,
            dir_okay=True,
            readable=True,
            path_type=Path,
        ),

        "help": (
            "Path to a directory with DICOMs "
            "for a single scanning session"
        ),
    },

    "sub-id": {

        "help": (
            "Pseudonymized TRR379 Q01 subject "
            "identifier corresponding to the scan"
        ),
    },

    "ses-id": {

        "help": (
            "TRR379 session identifier "
            "for this subject"
        ),
    },

    ("-m", "--message"): {

        "help": (
            "Commit message WHY a registration "
            "was updated"
        ),
    },

    "--container-cache-dir": {

        "envvar":
            f"{envvar_prefix}CONTAINERCACHE_DIR",

        "show_envvar":
            True,

        "type": click.Path(
            exists=True,
            file_okay=False,
            dir_okay=True,
            path_type=Path,
        ),

        "help": (
            "Directory for locally caching "
            "container image files"
        ),
    },
}


@cli.command(
    "create-session-dataset",
    help=(
        "Create and deposit a new DataLad dataset "
        "for a single scanning session"
    ),
)
@click.argument(
    "dicoms",
    **param_cfg["dicoms"],
)
@click.argument(
    "sub-id",
    **param_cfg["sub-id"],
)
@click.argument(
    "ses-id",
    **param_cfg["ses-id"],
)
def create_session_dataset(
    dicoms,
    sub_id,
    ses_id,
):

    try:

        with tempfile.TemporaryDirectory() as wdir:

            gitsha, deposit_url, dsid = (
                _do_session_dataset(
                    Path(wdir),
                    dicoms=dicoms,
                    sub_id=sub_id,
                    ses_id=ses_id,
                )
            )

    except ValueError as e:

        post_failure(
            str(e)
        )

        sys.exit(1)

    post_success(
        "DICOM session DataLad dataset "
        f"deposited at {deposit_url}"
    )


def _do_session_dataset(
    wdir: Path,
    dicoms: Path,
    sub_id: str,
    ses_id: str,
):

    deposit_url = (
        get_session_dataset_url(
            sub_id,
            ses_id,
        )
    )

    if call_git_success(
        [
            "ls-remote",
            deposit_url,
        ],
        capture_output=True,
    ):

        msg = (
            "There is already a DICOM session "
            f"dataset for {sub_id}-{ses_id} at "
            f"{deposit_url}. "
            "If the subject/session identifier "
            "is correct, this dataset needs to "
            "be updated manually if needed. "
            "DO NOT DELETE this existing repository, "
            "because this would invalidate existing "
            "metadata."
        )

        raise ValueError(msg)

    wt = create_dataset(
        wdir,
        private_annex=True,
    )

    archive_dest_path = (
        wt.path
        / get_archive_dest_rpath(
            dicoms,
            sub_id,
            ses_id,
        )
    )

    make_dicom_archive(
        dicoms,
        archive_dest_path,
        track,
    )

    call_git(
        [
            "annex",
            "add",
            archive_dest_path,
        ],
        cwd=wt.path,
    )

    call_git(
        [
            "commit",
            "-m",
            "ingest DICOM session archive",
        ],
        cwd=wt.path,
    )

    call_git(
        [
            "remote",
            "add",
            "origin",
            deposit_url,
        ],
        cwd=wt.path,
    )

    call_git(
        [
            "push",
            "--all",
        ],
        cwd=wt.path,
    )

    gitsha = call_git_oneline(
        [
            "rev-parse",
            "HEAD",
        ],
        cwd=wt.path,
    )

    return (
        gitsha,
        deposit_url,
        wt.config.get(
            "datalad.dataset.id"
        ).value,
    )


@cli.command(
    "register-session-dataset",
    help=(
        "Register a single session dataset in "
        "the all scan sessions DataLad dataset"
    ),
)
@click.argument(
    "sub-id",
    **param_cfg["sub-id"],
)
@click.argument(
    "ses-id",
    **param_cfg["ses-id"],
)
@click.option(
    "-m",
    "--message",
    **param_cfg[
        ("-m", "--message")
    ],
)
def register_session_dataset(
    sub_id,
    ses_id,
    message: str | None = None,
):

    deposit_url = (
        get_session_dataset_url(
            sub_id,
            ses_id,
        )
    )

    dsid, gitsha = (
        get_remote_dataset_state(
            deposit_url
        )
    )

    with tempfile.TemporaryDirectory() as wdir:

        gitsha, updated = (
            _do_dicoms_dataset(
                Path(wdir),
                sub_id=sub_id,
                ses_id=ses_id,
                gitsha=gitsha,
                deposit_url=deposit_url,
                dsid=dsid,
                message=message,
            )
        )

    if updated:

        post_success(
            "Pushed updated DICOM sessions "
            "DataLad dataset"
        )

    else:

        post_noop_result(
            "Registration found up to date, "
            "no changes made"
        )


def _do_dicoms_dataset(
    wdir: Path,
    sub_id: str,
    ses_id: str,
    gitsha: str,
    deposit_url: str,
    dsid: str,
    message: str | None,
) -> tuple[str, bool]:

    call_git(
        [
            "clone",
            "-q",
            "--depth",
            "1",
            all_dicoms_ds_url,
            wdir,
        ]
    )

    subm_rpath = (
        get_session_dataset_rpath(
            sub_id,
            ses_id,
        )
    )

    call_git(
        [
            "update-index",
            "--add",
            "--replace",
            "--cacheinfo",
            "160000",
            gitsha,
            subm_rpath,
        ],
        cwd=wdir,
    )

    for k, v in (
        (
            "path",
            subm_rpath,
        ),
        (
            "url",
            deposit_url,
        ),
        (
            "datalad-id",
            dsid,
        ),
    ):

        call_git(
            [
                "config",
                "--file",
                ".gitmodules",
                f"submodule.{subm_rpath}.{k}",
                v,
            ],
            cwd=wdir,
        )

    call_git(
        [
            "add",
            ".gitmodules",
        ],
        cwd=wdir,
    )

    updated = False

    if has_staged_changes(wdir):

        call_git(
            [
                "commit",
                "-m",
                message
                or (
                    f"add DICOM "
                    f"{sub_id}-{ses_id} "
                    "session dataset"
                ),
            ],
            cwd=wdir,
        )

        call_git(
            ["push"],
            cwd=wdir,
        )

        updated = True

    gitsha = call_git_oneline(
        [
            "rev-parse",
            "HEAD",
        ],
        cwd=wdir,
    )

    return (
        gitsha,
        updated,
    )


@cli.command(
    "update-bids-sourcedata",
    help=(
        "Update the BIDS dataset to the latest "
        "DICOM source data"
    ),
)
def update_bids_sourcedata():

    with tempfile.TemporaryDirectory() as wdir:

        updated = (
            _do_update_bids_sourcedata(
                Path(wdir)
            )
        )

    if updated:

        post_success(
            "Pushed updated BIDS DataLad dataset"
        )

    else:

        post_noop_result(
            "DICOM dataset registration found "
            "up to date, no changes made"
        )


def _do_update_bids_sourcedata(
    wdir: Path,
    gitsha: str | None = None,
) -> bool:

    if gitsha is None:

        _, gitsha = (
            get_remote_dataset_state(
                all_dicoms_ds_url
            )
        )

    call_git(
        [
            "clone",
            "-q",
            "--depth",
            "1",
            bids_ds_url,
            wdir,
        ]
    )

    call_git(
        [
            "update-index",
            "--add",
            "--replace",
            "--cacheinfo",
            "160000",
            gitsha,
            "sourcedata/dicoms",
        ],
        cwd=wdir,
    )

    if has_staged_changes(wdir):

        call_git(
            [
                "commit",
                "-m",
                "update DICOM source dataset",
            ],
            cwd=wdir,
        )

        call_git(
            ["push"],
            cwd=wdir,
        )

        return True

    return False


def _init_container_from_cache(
    wdir: Path,
    container_cache_dir: Path,
):

    container_ds_rpath = (
        "code/heudiconv"
    )

    shallow_clone_submodule(
        wdir,
        container_ds_rpath,
    )

    container_ds_path = (
        wdir
        / container_ds_rpath
    )

    call_git(
        [
            "fetch",
            "-q",
            "origin",
            "git-annex:git-annex",
            "--depth",
            "1",
        ],
        cwd=container_ds_path,
    )

    call_git(
        [
            "annex",
            "init",
        ],
        cwd=container_ds_path,
    )

    container_img_rpath = (
        call_git_oneline(
            [
                "config",
                "--file",
                ".datalad/config",
                "datalad.containers.apptainer.image",
            ],
            cwd=container_ds_path,
        )
    )

    img_key = (
        call_git_oneline(
            [
                "annex",
                "lookupkey",
                container_img_rpath,
            ],
            cwd=container_ds_path,
        )
    )

    cached_img = (
        container_cache_dir
        / img_key
    )

    if cached_img.exists():

        with Progress() as progress:

            progress.add_task(
                "Injecting cached container image",
                total=None,
            )

        copyfile(
            cached_img,
            container_ds_path
            / ".img_key",
        )

        call_git(
            [
                "annex",
                "reinject",
                ".img_key",
                container_img_rpath,
            ],
            cwd=container_ds_path,
        )


def _cache_container(
    wdir: Path,
    container_cache_dir: Path,
):

    container_img_rpath = (
        call_git_oneline(
            [
                "config",
                "--file",
                ".datalad/config",
                "datalad.containers.apptainer.image",
            ],
            cwd=wdir,
        )
    )

    img_key = (
        call_git_oneline(
            [
                "annex",
                "lookupkey",
                container_img_rpath,
            ],
            cwd=wdir,
        )
    )

    key_rpath = PurePosixPath(
        call_git_oneline(
            [
                "annex",
                "contentlocation",
                img_key,
            ],
            cwd=wdir,
        )
    )

    if not (
        container_cache_dir
        / key_rpath.name
    ).exists():

        with Progress() as progress:

            progress.add_task(
                "Cache container image",
                total=None,
            )

        copyfile(
            wdir
            / key_rpath,
            container_cache_dir
            / key_rpath.name,
        )


def _init_annex(
    repodir: Path,
    *,
    private,
):

    if private:

        call_git(
            [
                "config",
                "annex.private",
                "true",
            ],
            cwd=repodir,
        )

    try:

        call_git(
            [
                "fetch",
                "-q",
                "origin",
                "git-annex:git-annex",
                "--depth",
                "1",
            ],
            cwd=repodir,
        )

    except CommandError:

        pass

    call_git(
        [
            "annex",
            "init",
        ],
        cwd=repodir,
    )



def _dataset_relative_path(
    dataset_root: Path,
    path: Path,
) -> str:
    return path.relative_to(dataset_root).as_posix()


def _unlock_metadata_for_edit(
    dataset_root: Path,
    session_dir: Path,
):
    import datalad.api as dl

    metadata_paths = []
    metadata_paths.extend(sorted(session_dir.glob("*_scans.tsv")))

    for dirname in ("func", "fmap", "dwi"):
        folder = session_dir / dirname
        if folder.exists():
            metadata_paths.extend(sorted(folder.glob("*.json")))

    relative_paths = [
        _dataset_relative_path(dataset_root, path)
        for path in metadata_paths
        if path.exists() or path.is_symlink()
    ]

    if not relative_paths:
        return

    print("\nUnlocking KJP metadata files for post conversion edits...")

    results = dl.unlock(
        dataset=dl.Dataset(dataset_root),
        path=relative_paths,
        result_renderer="disabled",
        on_failure="ignore",
    )

    failures = [
        result for result in results
        if result.get("status") in {"error", "impossible"}
    ]

    if failures:
        details = "\n".join(
            f"  {result.get('path')}: {result.get('message')}"
            for result in failures
        )
        raise RuntimeError(
            "DataLad could not unlock one or more metadata files:\n"
            f"{details}"
        )


def _remove_tracked_file(
    dataset_root: Path,
    path: Path,
):
    if not path.exists() and not path.is_symlink():
        return

    relative_path = _dataset_relative_path(dataset_root, path)

    if call_git_success(
        ["ls-files", "--error-unmatch", "--", relative_path],
        cwd=dataset_root,
        capture_output=True,
    ):
        call_git(
            ["rm", "-f", "--", relative_path],
            cwd=dataset_root,
        )
    else:
        path.unlink(missing_ok=True)


def _rename_tracked_file(
    dataset_root: Path,
    source: Path,
    destination: Path,
):
    if source == destination:
        return

    if not source.exists() and not source.is_symlink():
        raise FileNotFoundError(f"Cannot rename missing file: {source}")

    if destination.exists() or destination.is_symlink():
        _remove_tracked_file(dataset_root, destination)

    source_relative = _dataset_relative_path(dataset_root, source)
    destination_relative = _dataset_relative_path(dataset_root, destination)

    if call_git_success(
        ["ls-files", "--error-unmatch", "--", source_relative],
        cwd=dataset_root,
        capture_output=True,
    ):
        call_git(
            ["mv", "-f", "--", source_relative, destination_relative],
            cwd=dataset_root,
        )
    else:
        source.rename(destination)


def _load_json(
    path: Path,
) -> dict:

    with path.open(
        "r",
        encoding="utf-8",
    ) as f:

        return json.load(f)


def _save_json(
    path: Path,
    data: dict,
):

    with path.open(
        "w",
        encoding="utf-8",
    ) as f:

        json.dump(
            data,
            f,
            indent=4,
            ensure_ascii=False,
        )

        f.write("\n")


def _paired_json(
    nii_path: Path,
) -> Path:

    name = nii_path.name

    if not name.endswith(
        ".nii.gz"
    ):

        raise ValueError(
            f"Expected .nii.gz file: "
            f"{nii_path}"
        )

    return nii_path.with_name(
        name[:-7]
        + ".json"
    )


def _nifti_shape(
    path: Path,
):

    return tuple(
        nib.load(
            str(path)
        ).shape
    )


def _cleanup_task_vid(
    dataset_root: Path,
    session_dir: Path,
    sub_id: str,
    ses_id: str,
):

    func_dir = session_dir / "func"

    if not func_dir.exists():
        print("WARNING: No func directory found.")
        return {"renamed_from": None, "renamed_to": None, "removed": []}

    prefix = f"sub-{sub_id}_ses-{ses_id}"
    final_nii = func_dir / f"{prefix}_task-vid_bold.nii.gz"
    candidates = sorted(func_dir.glob(f"{prefix}_task-vid_bold*.nii.gz"))

    print("\nTask vid BOLD candidates:")

    if not candidates:
        print("  None found.")
        return {"renamed_from": None, "renamed_to": None, "removed": []}

    candidate_info = []
    for path in candidates:
        shape = _nifti_shape(path)
        print(f"  {path.name}: {shape}")
        if len(shape) == 4 and shape[3] > 1:
            candidate_info.append((shape[3], path.stat().st_size, path, shape))

    if not candidate_info:
        raise RuntimeError("No valid 4D task vid BOLD candidate was found.")

    candidate_info.sort(key=lambda x: (x[0], x[1]))
    _, _, keep_nii, keep_shape = candidate_info[-1]

    print("\nKeeping task vid BOLD:")
    print(f"  {keep_nii.name}")
    print(f"  shape: {keep_shape}")

    keep_json = _paired_json(keep_nii)
    old_keep_relative = keep_nii.relative_to(session_dir).as_posix()
    removed_relative = []

    for path in candidates:
        if path == keep_nii:
            continue

        removed_relative.append(path.relative_to(session_dir).as_posix())
        json_path = _paired_json(path)
        print(f"Removing: {path.name}")
        _remove_tracked_file(dataset_root, path)
        _remove_tracked_file(dataset_root, json_path)

    final_json = _paired_json(final_nii)

    if keep_nii != final_nii:
        print("Renaming task vid BOLD:")
        print(f"  {keep_nii.name}")
        print(f"  -> {final_nii.name}")
        _rename_tracked_file(dataset_root, keep_nii, final_nii)

        if keep_json.exists() or keep_json.is_symlink():
            _rename_tracked_file(dataset_root, keep_json, final_json)

    return {
        "renamed_from": old_keep_relative,
        "renamed_to": final_nii.relative_to(session_dir).as_posix(),
        "removed": removed_relative,
    }

def _update_scans_tsv(
    session_dir: Path,
    sub_id: str,
    ses_id: str,
    cleanup_result: dict,
):

    scans_tsv = (
        session_dir
        / (
            f"sub-{sub_id}_"
            f"ses-{ses_id}_"
            "scans.tsv"
        )
    )

    if not scans_tsv.exists():

        print(
            "WARNING: scans.tsv not found:"
        )

        print(
            f"  {scans_tsv}"
        )

        return

    renamed_from = (
        cleanup_result.get(
            "renamed_from"
        )
    )

    renamed_to = (
        cleanup_result.get(
            "renamed_to"
        )
    )

    removed = set(
        cleanup_result.get(
            "removed",
            [],
        )
    )

    with scans_tsv.open(
        "r",
        encoding="utf-8",
        newline="",
    ) as f:

        reader = csv.DictReader(
            f,
            delimiter="\t",
        )

        fieldnames = (
            reader.fieldnames
        )

        rows = list(
            reader
        )

    if (
        fieldnames is None
        or "filename"
        not in fieldnames
    ):

        print(
            "WARNING: scans.tsv has no "
            "filename column."
        )

        return

    new_rows = []

    for row in rows:

        filename = (
            row["filename"]
        )

        if filename in removed:

            print(
                "Removing scans.tsv entry:"
            )

            print(
                f"  {filename}"
            )

            continue

        if (
            renamed_from
            and renamed_to
            and filename == renamed_from
        ):

            print(
                "Updating scans.tsv entry:"
            )

            print(
                f"  {renamed_from}"
            )

            print(
                f"  -> {renamed_to}"
            )

            row["filename"] = (
                renamed_to
            )

        new_rows.append(
            row
        )

    with scans_tsv.open(
        "w",
        encoding="utf-8",
        newline="",
    ) as f:

        writer = csv.DictWriter(
            f,
            fieldnames=fieldnames,
            delimiter="\t",
            lineterminator="\n",
        )

        writer.writeheader()

        writer.writerows(
            new_rows
        )


def _clear_b0_metadata(
    json_path: Path,
):

    if not json_path.exists():

        return False

    data = (
        _load_json(
            json_path
        )
    )

    changed = False

    for key in (
        "B0FieldIdentifier",
        "B0FieldSource",
    ):

        if key in data:

            del data[key]

            changed = True

    if changed:

        _save_json(
            json_path,
            data,
        )

    return changed


def _set_intended_for(
    fmap_json: Path,
    target_nii: Path,
    session_dir: Path,
    ses_id: str,
):

    if not fmap_json.exists():

        return

    data = (
        _load_json(
            fmap_json
        )
    )

    data.pop(
        "B0FieldIdentifier",
        None,
    )

    data.pop(
        "B0FieldSource",
        None,
    )

    if target_nii.exists():

        relative_target = (
            target_nii.relative_to(
                session_dir
            ).as_posix()
        )

        intended_for = (
            f"ses-{ses_id}/"
            f"{relative_target}"
        )

        data["IntendedFor"] = [
            intended_for
        ]

    else:

        data.pop(
            "IntendedFor",
            None,
        )

        print(
            "WARNING: IntendedFor target "
            "does not exist:"
        )

        print(
            f"  {target_nii}"
        )

    _save_json(
        fmap_json,
        data,
    )


def _finalize_kjp_bids_session(
    dataset_root: Path,
    sub_id: str,
    ses_id: str,
):

    session_dir = (
        dataset_root
        / f"sub-{sub_id}"
        / f"ses-{ses_id}"
    )

    if not session_dir.exists():

        raise RuntimeError(
            "Converted BIDS session "
            "directory does not exist: "
            f"{session_dir}"
        )

    print(
        "\nFinalizing Mannheim KJP "
        "BIDS session:"
    )

    print(
        f"  {session_dir}"
    )

    prefix = (
        f"sub-{sub_id}_"
        f"ses-{ses_id}"
    )

    _unlock_metadata_for_edit(
        dataset_root,
        session_dir,
    )

    cleanup_result = (
        _cleanup_task_vid(
            dataset_root,
            session_dir,
            sub_id,
            ses_id,
        )
    )

    _update_scans_tsv(
        session_dir,
        sub_id,
        ses_id,
        cleanup_result,
    )

    func_dir = (
        session_dir
        / "func"
    )

    fmap_dir = (
        session_dir
        / "fmap"
    )

    dwi_dir = (
        session_dir
        / "dwi"
    )

    rest_bold = (
        func_dir
        / (
            f"{prefix}_"
            "task-rest_bold.nii.gz"
        )
    )

    vid_bold = (
        func_dir
        / (
            f"{prefix}_"
            "task-vid_bold.nii.gz"
        )
    )

    vid_sbref_json = (
        func_dir
        / (
            f"{prefix}_"
            "task-vid_sbref.json"
        )
    )

    b1200_dwi = (
        dwi_dir
        / (
            f"{prefix}_"
            "acq-b1200_"
            "dir-PA_dwi.nii.gz"
        )
    )

    mshell_dwi = (
        dwi_dir
        / (
            f"{prefix}_"
            "acq-mshell_"
            "dir-PA_dwi.nii.gz"
        )
    )

    rest_fmaps = [

        fmap_dir
        / (
            f"{prefix}_"
            "acq-rest_"
            "dir-ap_epi.json"
        ),

        fmap_dir
        / (
            f"{prefix}_"
            "acq-rest_"
            "dir-pa_epi.json"
        ),
    ]

    vid_fmaps = [

        fmap_dir
        / (
            f"{prefix}_"
            "acq-vid_"
            "dir-ap_epi.json"
        ),

        fmap_dir
        / (
            f"{prefix}_"
            "acq-vid_"
            "dir-pa_epi.json"
        ),
    ]

    b1200_fmap = (
        fmap_dir
        / (
            f"{prefix}_"
            "acq-b1200_"
            "dir-AP_epi.json"
        )
    )

    mshell_fmap = (
        fmap_dir
        / (
            f"{prefix}_"
            "acq-mshell_"
            "dir-AP_epi.json"
        )
    )

    for fmap_json in rest_fmaps:

        _set_intended_for(
            fmap_json,
            rest_bold,
            session_dir,
            ses_id,
        )

    for fmap_json in vid_fmaps:

        _set_intended_for(
            fmap_json,
            vid_bold,
            session_dir,
            ses_id,
        )

    _set_intended_for(
        b1200_fmap,
        b1200_dwi,
        session_dir,
        ses_id,
    )

    _set_intended_for(
        mshell_fmap,
        mshell_dwi,
        session_dir,
        ses_id,
    )

    if vid_sbref_json.exists():

        sbref_data = (
            _load_json(
                vid_sbref_json
            )
        )

        sbref_data.pop(
            "IntendedFor",
            None,
        )

        sbref_data.pop(
            "B0FieldIdentifier",
            None,
        )

        sbref_data.pop(
            "B0FieldSource",
            None,
        )

        _save_json(
            vid_sbref_json,
            sbref_data,
        )

    for json_path in (
        func_dir.glob(
            "*.json"
        )
    ):

        _clear_b0_metadata(
            json_path
        )

    for json_path in (
        dwi_dir.glob(
            "*.json"
        )
    ):

        _clear_b0_metadata(
            json_path
        )

    print(
        "\nMannheim KJP finalization "
        "completed."
    )


@cli.command(
    "bids-conversion",
    help=(
        "Convert Mannheim KJP MRI data from "
        "a single DICOM session to BIDS format "
        "and incorporate it into the Q01 BIDS "
        "DataLad dataset."
    ),
)
@click.argument(
    "dicoms",
    **param_cfg["dicoms"],
)
@click.argument(
    "sub-id",
    **param_cfg["sub-id"],
)
@click.argument(
    "ses-id",
    **param_cfg["ses-id"],
)
@click.option(
    "-m",
    "--message",
    **param_cfg[
        ("-m", "--message")
    ],
)
@click.option(
    "--container-cache-dir",
    **param_cfg[
        "--container-cache-dir"
    ],
)
def bids_conversion(
    dicoms: Path,
    sub_id: str,
    ses_id: str,
    message: str | None = None,
    container_cache_dir:
        Path | None = None,
):

    with tempfile.TemporaryDirectory() as wdir:

        _do_bids_dataset(
            Path(wdir),
            dicoms,
            sub_id,
            ses_id,
            gitsha=None,
            container_cache_dir=(
                container_cache_dir
            ),
        )


def _do_bids_dataset(
    wdir: Path,
    dicoms: Path,
    sub_id: str,
    ses_id: str,
    gitsha: str | None = None,
    container_cache_dir:
        Path | None = None,
):

    _do_update_bids_sourcedata(
        Path(wdir),
        gitsha=gitsha,
    )

    shallow_clone_submodule(
        wdir,
        ".heudiconv",
    )

    call_git(
        [
            "switch",
            "main",
        ],
        cwd=(
            wdir
            / ".heudiconv"
        ),
    )

    for repodir in (
        wdir,
        wdir
        / ".heudiconv",
    ):

        _init_annex(
            repodir,
            private=True,
        )

    try:

        call_git_oneline(
            [
                "status",
                "--porcelain",
                ".heudiconv",
            ],
            cwd=wdir,
        )

        call_git(
            [
                "commit",
                "-m",
                (
                    "Sync external updates "
                    "to heudiconv log submodule"
                ),
                "--",
                ".heudiconv",
            ],
            cwd=wdir,
        )

    except AssertionError:

        pass


    dicom_src_rpath = (
        "sourcedata/dicoms"
    )

    shallow_clone_submodule(
        wdir,
        dicom_src_rpath,
    )

    dicom_session_rpath = (
        get_session_dataset_rpath(
            sub_id,
            ses_id,
        )
    )

    shallow_clone_submodule(
        wdir
        / dicom_src_rpath,
        dicom_session_rpath,
    )

    session_dicom_ds_path = (
        wdir
        / dicom_src_rpath
        / dicom_session_rpath
    )

    call_git(
        [
            "annex",
            "init",
        ],
        cwd=session_dicom_ds_path,
    )

    with Progress() as progress:

        progress.add_task(
            "Injecting DICOM session archive",
            total=None,
        )

        make_dicom_archive(
            dicoms,
            session_dicom_ds_path
            / ".archive_key",
            track,
        )

        call_git(
            [
                "annex",
                "reinject",
                ".archive_key",
                get_archive_dest_rpath(
                    dicoms,
                    sub_id,
                    ses_id,
                ),
            ],
            cwd=session_dicom_ds_path,
        )


    if container_cache_dir:

        _init_container_from_cache(
            wdir,
            container_cache_dir,
        )

    else:

        shallow_clone_submodule(
            wdir,
            "code/heudiconv",
        )

        subm_path = (
            wdir
            / "code"
            / "heudiconv"
        )

        call_git(
            [
                "fetch",
                "-q",
                "origin",
                "git-annex:git-annex",
                "--depth",
                "1",
            ],
            cwd=subm_path,
        )

        call_git(
            [
                "annex",
                "init",
            ],
            cwd=subm_path,
        )


    try:

        call_git_oneline(
            [
                "status",
                "--porcelain",
                "code/heudiconv",
            ],
            cwd=wdir,
        )

        post_failure(
            "Unincorporated changes to the "
            "heudiconv pipeline detected. "
            "Refusing to continue. "
            "Please check and advance the "
            "code/heudiconv submodule."
        )

        sys.exit(1)

    except AssertionError:

        pass


    pre_gitsha = (
        call_git_oneline(
            [
                "rev-parse",
                "HEAD",
            ],
            cwd=wdir,
        )
    )


    import datalad.api as dl


    with Progress() as progress:

        progress.add_task(
            "KJP BIDS conversion",
            total=None,
        )

        try:

            dl.containers_run(

                dataset=dl.Dataset(
                    wdir
                ),

                message=(
                    f"Convert KJP "
                    f"{sub_id}-{ses_id} "
                    "DICOM data"
                ),

                container_name=(
                    "code/heudiconv/"
                    "apptainer"
                ),

                outputs=[

                    (
                        f"sub-{sub_id}/"
                        f"ses-{ses_id}/"
                        f"sub-{sub_id}_"
                        f"ses-{ses_id}_"
                        "scans.tsv"
                    ),

                    (
                        f".heudiconv/"
                        f"sub-{sub_id}/"
                        f"ses-{ses_id}"
                    ),
                ],

                inputs=[

                    (
                        "sourcedata/dicoms/"
                        "code/"
                        "heuristic-q01-kjp.py"
                    ),

                    (
                        f"{dicom_src_rpath}/"
                        f"{dicom_session_rpath}"
                    ),
                ],

                cmd=(
                    "--bids notop "
                    "--overwrite "
                    "--minmeta "
                    "-o . "
                    "-f '{inputs[0]}' "
                    f'-s "{sub_id}" '
                    f'-ss "{ses_id}" '
                    "--files "
                    f"{dicom_src_rpath}/"
                    f"{dicom_session_rpath}"
                ),
            )

        except Exception:

            call_git(
                ["status"],
                cwd=wdir,
            )

            raise


    print(
        "\nRunning Mannheim KJP "
        "post conversion finalization..."
    )

    _finalize_kjp_bids_session(
        wdir,
        sub_id,
        ses_id,
    )


    print(
        "\nSaving Mannheim KJP post conversion changes with DataLad..."
    )

    dl.save(
        dataset=dl.Dataset(
            wdir
        ),

        path=[
            (
                f"sub-{sub_id}/"
                f"ses-{ses_id}"
            ),
        ],

        message=(
            "Finalize Mannheim KJP "
            f"BIDS session "
            f"{sub_id}-{ses_id}"
        ),
    )


    for repodir in (
        wdir,
        wdir
        / ".heudiconv",
    ):

        call_git(
            [
                "annex",
                "copy",
                "-t",
                "origin",
                ".",
            ],
            cwd=repodir,
        )

        call_git(
            ["push"],
            cwd=repodir,
        )


    if container_cache_dir:

        _cache_container(
            wdir
            / "code"
            / "heudiconv",
            container_cache_dir,
        )


    if pre_gitsha == call_git_oneline(
        [
            "rev-parse",
            "HEAD",
        ],
        cwd=wdir,
    ):

        post_noop_result(
            "No changes to tracked content."
        )

    else:

        post_success(
            "KJP conversion successful."
        )


@cli.command(
    "ingest-new-session",
    help=(
        "Perform all processing steps through "
        "BIDS conversion for a new Mannheim "
        "KJP scan."
    ),
)
@click.argument(
    "dicoms",
    **param_cfg["dicoms"],
)
@click.argument(
    "sub-id",
    **param_cfg["sub-id"],
)
@click.argument(
    "ses-id",
    **param_cfg["ses-id"],
)
@click.option(
    "--container-cache-dir",
    **param_cfg[
        "--container-cache-dir"
    ],
)
@click.pass_context
def ingest_new_session(
    ctx,
    dicoms: Path,
    sub_id: str,
    ses_id: str,
    container_cache_dir:
        Path | None = None,
) -> None:

    ctx.invoke(
        create_session_dataset,
        dicoms=dicoms,
        sub_id=sub_id,
        ses_id=ses_id,
    )

    ctx.invoke(
        register_session_dataset,
        sub_id=sub_id,
        ses_id=ses_id,
    )

    ctx.invoke(
        bids_conversion,
        dicoms=dicoms,
        sub_id=sub_id,
        ses_id=ses_id,
        container_cache_dir=(
            container_cache_dir
        ),
    )


if __name__ == "__main__":
    cli()