"""
Tests for t2smap.
"""

import os.path as op
from shutil import rmtree

import nibabel as nib

from tedana import workflows
from tedana.tests.utils import get_test_data_path


class TestT2smap():
    def test_basic_t2smap1(self):
        """
        A very simple test, to confirm that t2smap creates output
        files.
        """
        data_dir = get_test_data_path()
        data = [op.join(data_dir, 'echo1.nii.gz'),
                op.join(data_dir, 'echo2.nii.gz'),
                op.join(data_dir, 'echo3.nii.gz')]
        out_dir = 'TED.echo1.t2smap'
        workflows.t2smap_workflow(data, [14.5, 38.5, 62.5], combmode='t2s',
                                  fitmode='all', out_dir=out_dir)

        # Check outputs
        assert op.isfile(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        img = nib.load(op.join(out_dir, 'T2starmap.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'S0map.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-full_T2starmap.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-full_S0map.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        assert len(img.shape) == 4

    def test_basic_t2smap2(self):
        """
        A very simple test, to confirm that t2smap creates output
        files when fitmode is set to ts.
        """
        data_dir = get_test_data_path()
        data = [op.join(data_dir, 'echo1.nii.gz'),
                op.join(data_dir, 'echo2.nii.gz'),
                op.join(data_dir, 'echo3.nii.gz')]
        out_dir = 'TED.echo1.t2smap'
        workflows.t2smap_workflow(data, [14.5, 38.5, 62.5], combmode='t2s',
                                  fitmode='ts', out_dir=out_dir)

        # Check outputs
        assert op.isfile(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        img = nib.load(op.join(out_dir, 'T2starmap.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'S0map.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'desc-full_T2starmap.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'desc-full_S0map.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        assert len(img.shape) == 4

    def test_basic_t2smap3(self):
        """
        A very simple test, to confirm that t2smap creates output
        files when combmode is set to 'paid'.
        """
        data_dir = get_test_data_path()
        data = [op.join(data_dir, 'echo1.nii.gz'),
                op.join(data_dir, 'echo2.nii.gz'),
                op.join(data_dir, 'echo3.nii.gz')]
        out_dir = 'TED.echo1.t2smap'
        workflows.t2smap_workflow(data, [14.5, 38.5, 62.5], combmode='paid',
                                  fitmode='all', out_dir=out_dir)

        # Check outputs
        assert op.isfile(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        img = nib.load(op.join(out_dir, 'T2starmap.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'S0map.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-full_T2starmap.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-full_S0map.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        assert len(img.shape) == 4

    def test_basic_t2smap4(self):
        """
        A very simple test, to confirm that t2smap creates output
        files when combmode is set to 'paid' and fitmode is set to 'ts'.
        """
        data_dir = get_test_data_path()
        data = [op.join(data_dir, 'echo1.nii.gz'),
                op.join(data_dir, 'echo2.nii.gz'),
                op.join(data_dir, 'echo3.nii.gz')]
        out_dir = 'TED.echo1.t2smap'
        workflows.t2smap_workflow(data, [14.5, 38.5, 62.5], combmode='paid',
                                  fitmode='ts', out_dir=out_dir)

        # Check outputs
        assert op.isfile(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        img = nib.load(op.join(out_dir, 'T2starmap.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'S0map.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'desc-full_T2starmap.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'desc-full_S0map.nii.gz'))
        assert len(img.shape) == 4
        img = nib.load(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        assert len(img.shape) == 4

    def test_t2smap_cli(self):
        """
        Run test_basic_t2smap1, but use the CLI method.
        """
        data_dir = get_test_data_path()
        data = [op.join(data_dir, 'echo1.nii.gz'),
                op.join(data_dir, 'echo2.nii.gz'),
                op.join(data_dir, 'echo3.nii.gz')]
        out_dir = 'TED.echo1.t2smap'
        args = (['-d'] + data +
                ['-e', '14.5', '38.5', '62.5', '--combmode', 't2s',
                 '--fitmode', 'all', '--out-dir', out_dir])
        workflows.t2smap._main(args)

        # Check outputs
        img = nib.load(op.join(out_dir, 'T2starmap.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'S0map.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-full_T2starmap.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-full_S0map.nii.gz'))
        assert len(img.shape) == 3
        img = nib.load(op.join(out_dir, 'desc-optcom_bold.nii.gz'))
        assert len(img.shape) == 4

    def teardown_method(self):
        # Clean up folders
        rmtree('TED.echo1.t2smap')
