# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
#
#   See COPYING file distributed along with the NiBabel package for the
#   copyright and license terms.
#
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##

from os.path import join as pjoin

import numpy as np

from .. import load, Nifti1Image
from .. import brikhead

import pytest
from numpy.testing import assert_array_equal
from ..testing import data_path, assert_data_similar

from .test_fileslice import slicer_samples

EXAMPLE_IMAGES = [
    dict(
        head=pjoin(data_path, 'example4d+orig.HEAD'),
        fname=pjoin(data_path, 'example4d+orig.BRIK.gz'),
        shape=(33, 41, 25, 3),
        dtype=np.int16,
        affine=np.array([[-3.0,0,0,49.5],
                         [0,-3.0,0,82.312],
                         [0,0,3.0,-52.3511],
                         [0,0,0,1.0]]),
        zooms=(3., 3., 3., 3.),
        data_summary=dict(
            min=0,
            max=13722,
            mean=4266.76024636),
        is_proxy=True,
        space='ORIG',
        labels=['#0', '#1', '#2'],
        scaling=None),
    dict(
        head=pjoin(data_path, 'scaled+tlrc.HEAD'),
        fname=pjoin(data_path, 'scaled+tlrc.BRIK'),
        shape=(47, 54, 43, 1.),
        dtype=np.int16,
        affine=np.array([[3.0,0,0,-66.],
                         [0,3.0,0,-87.],
                         [0,0,3.0,-54.],
                         [0,0,0,1.0]]),
        zooms=(3., 3., 3., 0.),
        data_summary=dict(
            min=1.9416814999999998e-07,
            max=0.0012724615542099998,
            mean=0.00023919645351876782),
        is_proxy=True,
        space='TLRC',
        labels=['#0'],
        scaling=np.array([  3.88336300e-08]),
    )
]

EXAMPLE_BAD_IMAGES = [
    dict(
        head=pjoin(data_path, 'bad_datatype+orig.HEAD'),
        err=brikhead.AFNIImageError
    ),
    dict(
        head=pjoin(data_path, 'bad_attribute+orig.HEAD'),
        err=brikhead.AFNIHeaderError
    )
]

class TestAFNIHeader(object):
    module = brikhead
    test_files = EXAMPLE_IMAGES

    def test_makehead(self):
        for tp in self.test_files:
            head1 = self.module.AFNIHeader.from_fileobj(tp['head'])
            head2 = self.module.AFNIHeader.from_header(head1)
            assert head1 == head2
            with pytest.raises(self.module.AFNIHeaderError):
                self.module.AFNIHeader.from_header(header=None)
            with pytest.raises(self.module.AFNIHeaderError):
                self.module.AFNIHeader.from_header(tp['fname'])


class TestAFNIImage(object):
    module = brikhead
    test_files = EXAMPLE_IMAGES

    def test_brikheadfile(self):
        for tp in self.test_files:
            brik = self.module.load(tp['fname'])
            assert brik.get_data_dtype().type == tp['dtype']
            assert brik.shape == tp['shape']
            assert brik.header.get_zooms() == tp['zooms']
            assert_array_equal(brik.affine, tp['affine'])
            assert brik.header.get_space() == tp['space']
            data = brik.get_fdata()
            assert data.shape == tp['shape']
            assert_array_equal(brik.dataobj.scaling, tp['scaling'])
            assert brik.header.get_volume_labels() == tp['labels']

    def test_load(self):
        # Check highest level load of brikhead works
        for tp in self.test_files:
            img = self.module.load(tp['head'])
            data = img.get_fdata()
            assert data.shape == tp['shape']
            # min, max, mean values
            assert_data_similar(data, tp)
            # check if file can be converted to nifti
            ni_img = Nifti1Image.from_image(img)
            assert_array_equal(ni_img.affine, tp['affine'])
            assert_array_equal(ni_img.get_fdata(), data)

    def test_array_proxy_slicing(self):
        # Test slicing of array proxy
        for tp in self.test_files:
            img = self.module.load(tp['fname'])
            arr = img.get_fdata()
            prox = img.dataobj
            assert prox.is_proxy
            for sliceobj in slicer_samples(img.shape):
                assert_array_equal(arr[sliceobj], prox[sliceobj])


class TestBadFiles(object):
    module = brikhead
    test_files = EXAMPLE_BAD_IMAGES

    def test_brikheadfile(self):
        for tp in self.test_files:
            with pytest.raises(tp['err']):
                self.module.load(tp['head'])


class TestBadVars(object):
    module = brikhead
    vars = ['type = badtype-attribute\nname = BRICK_TYPES\ncount = 1\n1\n',
            'type = integer-attribute\ncount = 1\n1\n']

    def test_unpack_var(self):
        for var in self.vars:
            with pytest.raises(self.module.AFNIHeaderError):
                self.module._unpack_var(var)
