# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
import os
import pytest

from ....utils.filemanip import which
from ....testing import example_data
from .. import (
    get_custom_path,
    RegAverage,
    RegResample,
    RegJacobian,
    RegTools,
    RegMeasure,
    RegTransform,
)


def no_nifty_tool(cmd=None):
    return which(cmd) is None


@pytest.mark.skipif(
    no_nifty_tool(cmd="reg_resample"),
    reason="niftyreg is not installed. reg_resample not found.",
)
def test_reg_resample_res():
    """tests for reg_resample interface"""
    # Create a reg_resample object
    nr_resample = RegResample()

    # Check if the command is properly defined
    assert nr_resample.cmd == get_custom_path("reg_resample")

    # test raising error with mandatory args absent
    with pytest.raises(ValueError):
        nr_resample.run()

    # Resample res
    ref_file = example_data("im1.nii")
    flo_file = example_data("im2.nii")
    trans_file = example_data("warpfield.nii")
    nr_resample.inputs.ref_file = ref_file
    nr_resample.inputs.flo_file = flo_file
    nr_resample.inputs.trans_file = trans_file
    nr_resample.inputs.inter_val = "LIN"
    nr_resample.inputs.omp_core_val = 4

    cmd_tmp = "{cmd} -flo {flo} -inter 1 -omp 4 -ref {ref} -trans {trans} \
-res {res}"

    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_resample"),
        flo=flo_file,
        ref=ref_file,
        trans=trans_file,
        res="im2_res.nii.gz",
    )

    assert nr_resample.cmdline == expected_cmd

    # test_reg_resample_blank()
    nr_resample_2 = RegResample(type="blank", inter_val="LIN", omp_core_val=4)
    ref_file = example_data("im1.nii")
    flo_file = example_data("im2.nii")
    trans_file = example_data("warpfield.nii")
    nr_resample_2.inputs.ref_file = ref_file
    nr_resample_2.inputs.flo_file = flo_file
    nr_resample_2.inputs.trans_file = trans_file

    cmd_tmp = "{cmd} -flo {flo} -inter 1 -omp 4 -ref {ref} -trans {trans} \
-blank {blank}"

    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_resample"),
        flo=flo_file,
        ref=ref_file,
        trans=trans_file,
        blank="im2_blank.nii.gz",
    )

    assert nr_resample_2.cmdline == expected_cmd


@pytest.mark.skipif(
    no_nifty_tool(cmd="reg_jacobian"),
    reason="niftyreg is not installed. reg_jacobian not found.",
)
def test_reg_jacobian_jac():
    """Test interface for RegJacobian"""
    # Create a reg_jacobian object
    nr_jacobian = RegJacobian()

    # Check if the command is properly defined
    assert nr_jacobian.cmd == get_custom_path("reg_jacobian")

    # test raising error with mandatory args absent
    with pytest.raises(ValueError):
        nr_jacobian.run()

    # Test Reg Jacobian: jac
    ref_file = example_data("im1.nii")
    trans_file = example_data("warpfield.nii")
    nr_jacobian.inputs.ref_file = ref_file
    nr_jacobian.inputs.trans_file = trans_file
    nr_jacobian.inputs.omp_core_val = 4

    cmd_tmp = "{cmd} -omp 4 -ref {ref} -trans {trans} -jac {jac}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_jacobian"),
        ref=ref_file,
        trans=trans_file,
        jac="warpfield_jac.nii.gz",
    )

    assert nr_jacobian.cmdline == expected_cmd

    # Test Reg Jacobian: jac m
    nr_jacobian_2 = RegJacobian(type="jacM", omp_core_val=4)
    ref_file = example_data("im1.nii")
    trans_file = example_data("warpfield.nii")
    nr_jacobian_2.inputs.ref_file = ref_file
    nr_jacobian_2.inputs.trans_file = trans_file

    cmd_tmp = "{cmd} -omp 4 -ref {ref} -trans {trans} -jacM {jac}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_jacobian"),
        ref=ref_file,
        trans=trans_file,
        jac="warpfield_jacM.nii.gz",
    )

    assert nr_jacobian_2.cmdline == expected_cmd

    # Test Reg Jacobian: jac l
    nr_jacobian_3 = RegJacobian(type="jacL", omp_core_val=4)
    ref_file = example_data("im1.nii")
    trans_file = example_data("warpfield.nii")
    nr_jacobian_3.inputs.ref_file = ref_file
    nr_jacobian_3.inputs.trans_file = trans_file

    cmd_tmp = "{cmd} -omp 4 -ref {ref} -trans {trans} -jacL {jac}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_jacobian"),
        ref=ref_file,
        trans=trans_file,
        jac="warpfield_jacL.nii.gz",
    )

    assert nr_jacobian_3.cmdline == expected_cmd


@pytest.mark.skipif(
    no_nifty_tool(cmd="reg_tools"),
    reason="niftyreg is not installed. reg_tools not found.",
)
def test_reg_tools_mul():
    """tests for reg_tools interface"""
    # Create a reg_tools object
    nr_tools = RegTools()

    # Check if the command is properly defined
    assert nr_tools.cmd == get_custom_path("reg_tools")

    # test raising error with mandatory args absent
    with pytest.raises(ValueError):
        nr_tools.run()

    # Test reg_tools: mul
    in_file = example_data("im1.nii")
    nr_tools.inputs.in_file = in_file
    nr_tools.inputs.mul_val = 4
    nr_tools.inputs.omp_core_val = 4

    cmd_tmp = "{cmd} -in {in_file} -mul 4.0 -omp 4 -out {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_tools"), in_file=in_file, out_file="im1_tools.nii.gz"
    )

    assert nr_tools.cmdline == expected_cmd

    # Test reg_tools: iso
    nr_tools_2 = RegTools(iso_flag=True, omp_core_val=4)
    in_file = example_data("im1.nii")
    nr_tools_2.inputs.in_file = in_file

    cmd_tmp = "{cmd} -in {in_file} -iso -omp 4 -out {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_tools"), in_file=in_file, out_file="im1_tools.nii.gz"
    )

    assert nr_tools_2.cmdline == expected_cmd


@pytest.mark.skipif(
    no_nifty_tool(cmd="reg_average"),
    reason="niftyreg is not installed. reg_average not found.",
)
def test_reg_average():
    """tests for reg_average interface"""
    # Create a reg_average object
    nr_average = RegAverage()

    # Check if the command is properly defined
    assert nr_average.cmd == get_custom_path("reg_average")

    # Average niis
    one_file = example_data("im1.nii")
    two_file = example_data("im2.nii")
    three_file = example_data("im3.nii")
    nr_average.inputs.avg_files = [one_file, two_file, three_file]
    nr_average.inputs.omp_core_val = 1
    generated_cmd = nr_average.cmdline

    # Read the reg_average_cmd
    reg_average_cmd = os.path.join(os.getcwd(), "reg_average_cmd")
    with open(reg_average_cmd, "rb") as f_obj:
        argv = f_obj.read()
    os.remove(reg_average_cmd)

    expected_argv = "%s %s -avg %s %s %s -omp 1" % (
        get_custom_path("reg_average"),
        os.path.join(os.getcwd(), "avg_out.nii.gz"),
        one_file,
        two_file,
        three_file,
    )

    assert argv.decode("utf-8") == expected_argv

    # Test command line with text file
    expected_cmd = "%s --cmd_file %s" % (
        get_custom_path("reg_average"),
        reg_average_cmd,
    )

    assert generated_cmd == expected_cmd

    # Test Reg Average: average txt
    nr_average_2 = RegAverage()
    one_file = example_data("TransformParameters.0.txt")
    two_file = example_data("ants_Affine.txt")
    three_file = example_data("elastix.txt")
    nr_average_2.inputs.avg_files = [one_file, two_file, three_file]
    nr_average_2.inputs.omp_core_val = 1
    generated_cmd = nr_average_2.cmdline

    # Read the reg_average_cmd
    reg_average_cmd = os.path.join(os.getcwd(), "reg_average_cmd")
    with open(reg_average_cmd, "rb") as f_obj:
        argv = f_obj.read()
    os.remove(reg_average_cmd)

    expected_argv = "%s %s -avg %s %s %s -omp 1" % (
        get_custom_path("reg_average"),
        os.path.join(os.getcwd(), "avg_out.txt"),
        one_file,
        two_file,
        three_file,
    )

    assert argv.decode("utf-8") == expected_argv

    # Test Reg Average: average list
    nr_average_3 = RegAverage()
    one_file = example_data("TransformParameters.0.txt")
    two_file = example_data("ants_Affine.txt")
    three_file = example_data("elastix.txt")
    nr_average_3.inputs.avg_lts_files = [one_file, two_file, three_file]
    nr_average_3.inputs.omp_core_val = 1
    generated_cmd = nr_average_3.cmdline

    # Read the reg_average_cmd
    reg_average_cmd = os.path.join(os.getcwd(), "reg_average_cmd")
    with open(reg_average_cmd, "rb") as f_obj:
        argv = f_obj.read()
    os.remove(reg_average_cmd)

    expected_argv = "%s %s -avg_lts %s %s %s -omp 1" % (
        get_custom_path("reg_average"),
        os.path.join(os.getcwd(), "avg_out.txt"),
        one_file,
        two_file,
        three_file,
    )

    assert argv.decode("utf-8") == expected_argv

    # Test Reg Average: average ref
    nr_average_4 = RegAverage()
    ref_file = example_data("anatomical.nii")
    one_file = example_data("im1.nii")
    two_file = example_data("im2.nii")
    three_file = example_data("im3.nii")
    trans1_file = example_data("roi01.nii")
    trans2_file = example_data("roi02.nii")
    trans3_file = example_data("roi03.nii")
    nr_average_4.inputs.warp_files = [
        trans1_file,
        one_file,
        trans2_file,
        two_file,
        trans3_file,
        three_file,
    ]
    nr_average_4.inputs.avg_ref_file = ref_file
    nr_average_4.inputs.omp_core_val = 1
    generated_cmd = nr_average_4.cmdline

    # Read the reg_average_cmd
    reg_average_cmd = os.path.join(os.getcwd(), "reg_average_cmd")
    with open(reg_average_cmd, "rb") as f_obj:
        argv = f_obj.read()
    os.remove(reg_average_cmd)

    expected_argv = "%s %s -avg_tran %s -omp 1 %s %s %s %s %s %s" % (
        get_custom_path("reg_average"),
        os.path.join(os.getcwd(), "avg_out.nii.gz"),
        ref_file,
        trans1_file,
        one_file,
        trans2_file,
        two_file,
        trans3_file,
        three_file,
    )

    assert argv.decode("utf-8") == expected_argv

    # Test Reg Average: demean3
    nr_average_5 = RegAverage()
    ref_file = example_data("anatomical.nii")
    one_file = example_data("im1.nii")
    two_file = example_data("im2.nii")
    three_file = example_data("im3.nii")
    aff1_file = example_data("TransformParameters.0.txt")
    aff2_file = example_data("ants_Affine.txt")
    aff3_file = example_data("elastix.txt")
    trans1_file = example_data("roi01.nii")
    trans2_file = example_data("roi02.nii")
    trans3_file = example_data("roi03.nii")
    nr_average_5.inputs.warp_files = [
        aff1_file,
        trans1_file,
        one_file,
        aff2_file,
        trans2_file,
        two_file,
        aff3_file,
        trans3_file,
        three_file,
    ]
    nr_average_5.inputs.demean3_ref_file = ref_file
    nr_average_5.inputs.omp_core_val = 1
    generated_cmd = nr_average_5.cmdline

    # Read the reg_average_cmd
    reg_average_cmd = os.path.join(os.getcwd(), "reg_average_cmd")
    with open(reg_average_cmd, "rb") as f_obj:
        argv = f_obj.read()
    os.remove(reg_average_cmd)

    expected_argv = "%s %s -demean3 %s -omp 1 %s %s %s %s %s %s %s %s %s" % (
        get_custom_path("reg_average"),
        os.path.join(os.getcwd(), "avg_out.nii.gz"),
        ref_file,
        aff1_file,
        trans1_file,
        one_file,
        aff2_file,
        trans2_file,
        two_file,
        aff3_file,
        trans3_file,
        three_file,
    )

    assert argv.decode("utf-8") == expected_argv


@pytest.mark.skipif(
    no_nifty_tool(cmd="reg_transform"),
    reason="niftyreg is not installed. reg_transform not found.",
)
def test_reg_transform_def():
    """tests for reg_transform interface"""
    # Create a reg_transform object
    nr_transform = RegTransform()

    # Check if the command is properly defined
    assert nr_transform.cmd == get_custom_path("reg_transform")

    # Assign some input data
    trans_file = example_data("warpfield.nii")
    nr_transform.inputs.def_input = trans_file
    nr_transform.inputs.omp_core_val = 4

    cmd_tmp = "{cmd} -omp 4 -def {trans_file} {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_transform"),
        trans_file=trans_file,
        out_file=os.path.join(os.getcwd(), "warpfield_trans.nii.gz"),
    )

    assert nr_transform.cmdline == expected_cmd

    # Test reg_transform: def ref
    nr_transform_2 = RegTransform(omp_core_val=4)
    ref_file = example_data("im1.nii")
    trans_file = example_data("warpfield.nii")
    nr_transform_2.inputs.ref1_file = ref_file
    nr_transform_2.inputs.def_input = trans_file

    cmd_tmp = "{cmd} -ref {ref_file} -omp 4 -def {trans_file} {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_transform"),
        ref_file=ref_file,
        trans_file=trans_file,
        out_file=os.path.join(os.getcwd(), "warpfield_trans.nii.gz"),
    )

    assert nr_transform_2.cmdline == expected_cmd

    # Test reg_transform: comp nii
    nr_transform_3 = RegTransform(omp_core_val=4)
    ref_file = example_data("im1.nii")
    trans_file = example_data("warpfield.nii")
    trans2_file = example_data("anatomical.nii")
    nr_transform_3.inputs.ref1_file = ref_file
    nr_transform_3.inputs.comp_input2 = trans2_file
    nr_transform_3.inputs.comp_input = trans_file

    cmd_tmp = "{cmd} -ref {ref_file} -omp 4 -comp {trans1} {trans2} {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_transform"),
        ref_file=ref_file,
        trans1=trans_file,
        trans2=trans2_file,
        out_file=os.path.join(os.getcwd(), "warpfield_trans.nii.gz"),
    )

    assert nr_transform_3.cmdline == expected_cmd

    # Test reg_transform: comp txt
    nr_transform_4 = RegTransform(omp_core_val=4)
    aff1_file = example_data("ants_Affine.txt")
    aff2_file = example_data("elastix.txt")
    nr_transform_4.inputs.comp_input2 = aff2_file
    nr_transform_4.inputs.comp_input = aff1_file

    cmd_tmp = "{cmd} -omp 4 -comp {aff1} {aff2} {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_transform"),
        aff1=aff1_file,
        aff2=aff2_file,
        out_file=os.path.join(os.getcwd(), "ants_Affine_trans.txt"),
    )

    assert nr_transform_4.cmdline == expected_cmd

    # Test reg_transform: comp
    nr_transform_5 = RegTransform(omp_core_val=4)
    trans_file = example_data("warpfield.nii")
    aff_file = example_data("elastix.txt")
    nr_transform_5.inputs.comp_input2 = trans_file
    nr_transform_5.inputs.comp_input = aff_file

    cmd_tmp = "{cmd} -omp 4 -comp {aff} {trans} {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_transform"),
        aff=aff_file,
        trans=trans_file,
        out_file=os.path.join(os.getcwd(), "elastix_trans.nii.gz"),
    )

    assert nr_transform_5.cmdline == expected_cmd

    # Test reg_transform: flirt
    nr_transform_6 = RegTransform(omp_core_val=4)
    aff_file = example_data("elastix.txt")
    ref_file = example_data("im1.nii")
    in_file = example_data("im2.nii")
    nr_transform_6.inputs.flirt_2_nr_input = (aff_file, ref_file, in_file)

    cmd_tmp = "{cmd} -omp 4 -flirtAff2NR {aff} {ref} {in_file} {out_file}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_transform"),
        aff=aff_file,
        ref=ref_file,
        in_file=in_file,
        out_file=os.path.join(os.getcwd(), "elastix_trans.txt"),
    )

    assert nr_transform_6.cmdline == expected_cmd


@pytest.mark.skipif(
    no_nifty_tool(cmd="reg_measure"),
    reason="niftyreg is not installed. reg_measure not found.",
)
def test_reg_measure():
    """tests for reg_measure interface"""
    # Create a reg_measure object
    nr_measure = RegMeasure()

    # Check if the command is properly defined
    assert nr_measure.cmd == get_custom_path("reg_measure")

    # test raising error with mandatory args absent
    with pytest.raises(ValueError):
        nr_measure.run()

    # Assign some input data
    ref_file = example_data("im1.nii")
    flo_file = example_data("im2.nii")
    nr_measure.inputs.ref_file = ref_file
    nr_measure.inputs.flo_file = flo_file
    nr_measure.inputs.measure_type = "lncc"
    nr_measure.inputs.omp_core_val = 4

    cmd_tmp = "{cmd} -flo {flo} -lncc -omp 4 -out {out} -ref {ref}"
    expected_cmd = cmd_tmp.format(
        cmd=get_custom_path("reg_measure"),
        flo=flo_file,
        out="im2_lncc.txt",
        ref=ref_file,
    )

    assert nr_measure.cmdline == expected_cmd