#!/usr/bin/env fslpython


import                             shlex
import subprocess           as     sp
import numpy                as     np
import nibabel              as     nib
from   fsl.data.image       import Image
import fsl.transform.affine as     affine


def applywarp(src, ref):
    cmd = f'applywarp --in={src} --ref={ref} --usesqform --out=out'
    sp.run(shlex.split(cmd), check=True)
    return Image('out')


def make_image(fname, data, sform, sform_code, qform, qform_code):
    hdr = nib.Nifti1Header()
    hdr.set_sform(sform, sform_code)
    hdr.set_qform(qform, qform_code)
    img = Image(data, header=hdr)
    img.save(fname)
    return img

# fsl/fugue!6
#
# Make sure that applywarp --usesqform
# will correctly use the sform or qform
#
# FSL affine precedence rules are as follows:
#
#  1. If sform_code != 0, use sform; else
#  2. If qform_code != 0, use qform; else
#  3. Use a scaling matrix with pixdims along the
#     diagonal
def test_applywarp_uses_sqform_correctly():

    srcdata = np.zeros((20, 20, 20))
    refdata = np.zeros((20, 20, 20))

    # Data from the two images
    # aligned via their affines -
    # src should be resampled so
    # as to be identical to ref
    srcdata[ 6: 9,  6: 9,  6: 9] = 1
    refdata[16:19, 16:19, 16:19] = 1

    eye    = np.eye(4)
    srcaff = affine.scaleOffsetXform(1, [ 5,  5,  5])
    refaff = affine.scaleOffsetXform(1, [-5, -5, -5])

    # aligned via sform
    src    = make_image('src', srcdata, srcaff, 2, eye, 1)
    ref    = make_image('ref', refdata, refaff, 2, eye, 1)
    result = applywarp('src', 'ref')
    assert np.all(np.isclose(result.data, ref.data))

    # aligned via sform again
    src    = make_image('src', srcdata, srcaff, 2, eye, 0)
    ref    = make_image('ref', refdata, refaff, 2, eye, 0)
    result = applywarp('src', 'ref')
    assert np.all(np.isclose(result.data, ref.data))

    # aligned via qform
    src    = make_image('src', srcdata, eye, 0, srcaff, 2)
    ref    = make_image('ref', refdata, eye, 0, refaff, 2)
    result = applywarp('src', 'ref')
    assert np.all(np.isclose(result.data, ref.data))

    # not aligned - should result in a scaling
    # matrix being used (an identify matrix here)
    src    = make_image('src', srcdata, eye, 0, srcaff, 0)
    ref    = make_image('ref', refdata, eye, 0, refaff, 0)
    result = applywarp('src', 'ref')
    assert not np.all(np.isclose(result.data, ref.data))
    assert     np.all(np.isclose(result.data, src.data))


if __name__ == '__main__':
    test_applywarp_uses_sqform_correctly()