#! /usr/bin/env python

"""
Split a 4D NIFTI file into separate volumes.

Dec 11, 2023
David Gobbi
"""

usage="""
usage: split_nifti file.nii.gz

If the file has more than 3 dimensions, then it will be split into
separate volumes.  For example, if there are 2 volumes, the output is:

    file-vol1.nii.gz
    file-vol2.nii.gz

If "file.bval" exists in the same directory (e.g. from dcm2niix), then
the bvals will be used to name the output:

    file-vol1-b1000.nii.gz
    file-vol2-b1000.nii.gz
    file-vol3-b0.nii.gz

The ordering will match the ordering in the .bval file.
"""

import argparse
import gzip
import struct
import sys
import os
import os.path

def split_nifti(filename):
    """Split a NIFTI file into -vol1, -vol2, etc.
    If there is a bval, then add include -b<bval>.
    """
    # length of NIFTI header
    header_size = 348
    if filename.endswith(".nii.gz"):
        file = gzip.open(filename, 'rb')
        basename = filename[0:-7]
    elif filename.endswith(".nii"):
        file = open(filename, 'rb')
        basename = filename[0:-4]
    else:
        sys.stderr.write("File doesn't appear to be NIFTI: %s\n" % (filename))

    header = file.read(header_size)
    size_check = struct.unpack('i', header[0:4])[0]
    if size_check != header_size:
        sys.stderr.write("File doesn't appear to be NIFTI: %s\n" % (filename))
        return None

    # get the dims and pixdims
    dims = struct.unpack('h'*8, header[40:56])
    pixdim = struct.unpack('f'*8, header[76:108])
    vox_offset = int(struct.unpack('f', header[108:112])[0])
    nvol = 1
    for i in dims[4:1+dims[0]]:
        nvol = nvol * i
    pad  = file.read(vox_offset - header_size)

    if nvol == 1:
        file.close()
        return

    image = file.read()
    file.close()

    volsize = int(len(image)/nvol)
    newheader = header[:40] + struct.pack('h'*8, 3, dims[1], dims[2], dims[3], 1, 1, 1, 1) + header[56:]

    bvals = []
    bvalfile = basename + ".bval"
    if os.path.exists(bvalfile):
        file = open(bvalfile, 'r')
        t = file.read()
        file.close()
        bvals = [float(x) for x in t.split()]

    for i in range(nvol):
        if len(bvals) == nvol:
            btext = "-b%g" % bvals[i]
        else:
            btext = ""
        fname = "%s-vol%d%s.nii.gz" % (basename, i+1, btext)
        if os.path.exists(fname):
            continue
        file = gzip.open(fname, 'wb')
        file.write(newheader)
        file.write(pad)
        file.write(image[volsize*i:volsize*(i+1)])
        file.close()

if __name__ == '__main__':

    if len(sys.argv) < 2 or sys.argv[1].startswith('-') or not os.path.exists(sys.argv[1]):
        print(usage)
        sys.exit(0)

    filename = sys.argv[1]
    split_nifti(filename)