# 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.
#
### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
""" Tests for arrayproxy module
"""
from __future__ import division, print_function, absolute_import

import warnings
import gzip
import contextlib

import pickle
from io import BytesIO
from ..tmpdirs import InTemporaryDirectory

import numpy as np

from ..arrayproxy import (ArrayProxy, is_proxy, reshape_dataobj)
from ..openers import ImageOpener
from ..nifti1 import Nifti1Header

import mock

from numpy.testing import assert_array_equal, assert_array_almost_equal
from nose.tools import (assert_true, assert_false, assert_equal,
                        assert_not_equal, assert_raises)
from nibabel.testing import memmap_after_ufunc

from .test_fileslice import slicer_samples
from .test_openers import patch_indexed_gzip


class FunkyHeader(object):

    def __init__(self, shape):
        self.shape = shape

    def get_data_shape(self):
        return self.shape[:]

    def get_data_dtype(self):
        return np.int32

    def get_data_offset(self):
        return 16

    def get_slope_inter(self):
        return 1.0, 0.0

    def copy(self):
        # Not needed when we remove header property
        return FunkyHeader(self.shape)


class CArrayProxy(ArrayProxy):
    # C array memory layout
    order = 'C'


def test_init():
    bio = BytesIO()
    shape = [2, 3, 4]
    dtype = np.int32
    arr = np.arange(24, dtype=dtype).reshape(shape)
    bio.seek(16)
    bio.write(arr.tostring(order='F'))
    hdr = FunkyHeader(shape)
    ap = ArrayProxy(bio, hdr)
    assert_true(ap.file_like is bio)
    assert_equal(ap.shape, shape)
    # shape should be read only
    assert_raises(AttributeError, setattr, ap, 'shape', shape)
    # Get the data
    assert_array_equal(np.asarray(ap), arr)
    # Check we can modify the original header without changing the ap version
    hdr.shape[0] = 6
    assert_not_equal(ap.shape, shape)
    # Data stays the same, also
    assert_array_equal(np.asarray(ap), arr)
    # C order also possible
    bio = BytesIO()
    bio.seek(16)
    bio.write(arr.tostring(order='C'))
    ap = CArrayProxy(bio, FunkyHeader((2, 3, 4)))
    assert_array_equal(np.asarray(ap), arr)
    # Illegal init
    assert_raises(TypeError, ArrayProxy, bio, object())


def test_tuplespec():
    bio = BytesIO()
    shape = [2, 3, 4]
    dtype = np.int32
    arr = np.arange(24, dtype=dtype).reshape(shape)
    bio.seek(16)
    bio.write(arr.tostring(order='F'))
    # Create equivalent header and tuple specs
    hdr = FunkyHeader(shape)
    tuple_spec = (hdr.get_data_shape(), hdr.get_data_dtype(),
                  hdr.get_data_offset(), 1., 0.)
    ap_header = ArrayProxy(bio, hdr)
    ap_tuple = ArrayProxy(bio, tuple_spec)
    # Header and tuple specs produce identical behavior
    for prop in ('shape', 'dtype', 'offset', 'slope', 'inter', 'is_proxy'):
        assert_equal(getattr(ap_header, prop), getattr(ap_tuple, prop))
    for method, args in (('get_unscaled', ()), ('__array__', ()),
                         ('__getitem__', ((0, 2, 1), ))
                         ):
        assert_array_equal(getattr(ap_header, method)(*args),
                           getattr(ap_tuple, method)(*args))
    # Tuple-defined ArrayProxies have no header to store
    with warnings.catch_warnings():
        assert_true(ap_tuple.header is None)
    # Partial tuples of length 2-4 are also valid
    for n in range(2, 5):
        ArrayProxy(bio, tuple_spec[:n])
    # Bad tuple lengths
    assert_raises(TypeError, ArrayProxy, bio, ())
    assert_raises(TypeError, ArrayProxy, bio, tuple_spec[:1])
    assert_raises(TypeError, ArrayProxy, bio, tuple_spec + ('error',))


def write_raw_data(arr, hdr, fileobj):
    hdr.set_data_shape(arr.shape)
    hdr.set_data_dtype(arr.dtype)
    fileobj.write(b'\x00' * hdr.get_data_offset())
    fileobj.write(arr.tostring(order='F'))


def test_nifti1_init():
    bio = BytesIO()
    shape = (2, 3, 4)
    hdr = Nifti1Header()
    arr = np.arange(24, dtype=np.int16).reshape(shape)
    write_raw_data(arr, hdr, bio)
    hdr.set_slope_inter(2, 10)
    ap = ArrayProxy(bio, hdr)
    assert_true(ap.file_like == bio)
    assert_equal(ap.shape, shape)
    # Check there has been a copy of the header
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        assert_false(ap.header is hdr)
    # Get the data
    assert_array_equal(np.asarray(ap), arr * 2.0 + 10)
    with InTemporaryDirectory():
        f = open('test.nii', 'wb')
        write_raw_data(arr, hdr, f)
        f.close()
        ap = ArrayProxy('test.nii', hdr)
        assert_true(ap.file_like == 'test.nii')
        assert_equal(ap.shape, shape)
        assert_array_equal(np.asarray(ap), arr * 2.0 + 10)


def test_proxy_slicing():
    shapes = (15, 16, 17)
    for n_dim in range(1, len(shapes) + 1):
        shape = shapes[:n_dim]
        arr = np.arange(np.prod(shape)).reshape(shape)
        for offset in (0, 20):
            hdr = Nifti1Header()
            hdr.set_data_offset(offset)
            hdr.set_data_dtype(arr.dtype)
            hdr.set_data_shape(shape)
            for order, klass in ('F', ArrayProxy), ('C', CArrayProxy):
                fobj = BytesIO()
                fobj.write(b'\0' * offset)
                fobj.write(arr.tostring(order=order))
                prox = klass(fobj, hdr)
                for sliceobj in slicer_samples(shape):
                    assert_array_equal(arr[sliceobj], prox[sliceobj])
    # Check slicing works with scaling
    hdr.set_slope_inter(2.0, 1.0)
    fobj = BytesIO()
    fobj.write(b'\0' * offset)
    fobj.write(arr.tostring(order='F'))
    prox = ArrayProxy(fobj, hdr)
    sliceobj = (None, slice(None), 1, -1)
    assert_array_equal(arr[sliceobj] * 2.0 + 1.0, prox[sliceobj])


def test_is_proxy():
    # Test is_proxy function
    hdr = FunkyHeader((2, 3, 4))
    bio = BytesIO()
    prox = ArrayProxy(bio, hdr)
    assert_true(is_proxy(prox))
    assert_false(is_proxy(bio))
    assert_false(is_proxy(hdr))
    assert_false(is_proxy(np.zeros((2, 3, 4))))

    class NP(object):
        is_proxy = False
    assert_false(is_proxy(NP()))


def test_reshape_dataobj():
    # Test function that reshapes using method if possible
    shape = (1, 2, 3, 4)
    hdr = FunkyHeader(shape)
    bio = BytesIO()
    prox = ArrayProxy(bio, hdr)
    arr = np.arange(np.prod(shape), dtype=prox.dtype).reshape(shape)
    bio.write(b'\x00' * prox.offset + arr.tostring(order='F'))
    assert_array_equal(prox, arr)
    assert_array_equal(reshape_dataobj(prox, (2, 3, 4)),
                       np.reshape(arr, (2, 3, 4)))
    assert_equal(prox.shape, shape)
    assert_equal(arr.shape, shape)
    assert_array_equal(reshape_dataobj(arr, (2, 3, 4)),
                       np.reshape(arr, (2, 3, 4)))
    assert_equal(arr.shape, shape)

    class ArrGiver(object):

        def __array__(self):
            return arr

    assert_array_equal(reshape_dataobj(ArrGiver(), (2, 3, 4)),
                       np.reshape(arr, (2, 3, 4)))
    assert_equal(arr.shape, shape)


def test_reshaped_is_proxy():
    shape = (1, 2, 3, 4)
    hdr = FunkyHeader(shape)
    bio = BytesIO()
    prox = ArrayProxy(bio, hdr)
    assert_true(isinstance(prox.reshape((2, 3, 4)), ArrayProxy))
    minus1 = prox.reshape((2, -1, 4))
    assert_true(isinstance(minus1, ArrayProxy))
    assert_equal(minus1.shape, (2, 3, 4))
    assert_raises(ValueError, prox.reshape, (-1, -1, 4))
    assert_raises(ValueError, prox.reshape, (2, 3, 5))
    assert_raises(ValueError, prox.reshape, (2, -1, 5))


def test_get_unscaled():
    # Test fetch of raw array
    class FunkyHeader2(FunkyHeader):

        def get_slope_inter(self):
            return 2.1, 3.14
    shape = (2, 3, 4)
    hdr = FunkyHeader2(shape)
    bio = BytesIO()
    # Check standard read works
    arr = np.arange(24, dtype=np.int32).reshape(shape, order='F')
    bio.write(b'\x00' * hdr.get_data_offset())
    bio.write(arr.tostring(order='F'))
    prox = ArrayProxy(bio, hdr)
    assert_array_almost_equal(np.array(prox), arr * 2.1 + 3.14)
    # Check unscaled read works
    assert_array_almost_equal(prox.get_unscaled(), arr)


def test_mmap():
    # Unscaled should return mmap from suitable file, this can be tuned
    hdr = FunkyHeader((2, 3, 4))
    check_mmap(hdr, hdr.get_data_offset(), ArrayProxy)


def check_mmap(hdr, offset, proxy_class,
               has_scaling=False,
               unscaled_is_view=True):
    """ Assert that array proxies return memory maps as expected

    Parameters
    ----------
    hdr : object
        Image header instance
    offset : int
        Offset in bytes of image data in file (that we will write)
    proxy_class : class
        Class of image array proxy to test
    has_scaling : {False, True}
        True if the `hdr` says to apply scaling to the output data, False
        otherwise.
    unscaled_is_view : {True, False}
        True if getting the unscaled data returns a view of the array.  If
        False, then type of returned array will depend on whether numpy has the
        old viral (< 1.12) memmap behavior (returns memmap) or the new behavior
        (returns ndarray).  See: https://github.com/numpy/numpy/pull/7406
    """
    shape = hdr.get_data_shape()
    arr = np.arange(np.prod(shape), dtype=hdr.get_data_dtype()).reshape(shape)
    fname = 'test.bin'
    # Whether unscaled array memory backed by memory map (regardless of what
    # numpy says).
    unscaled_really_mmap = unscaled_is_view
    # Whether scaled array memory backed by memory map (regardless of what
    # numpy says).
    scaled_really_mmap = unscaled_really_mmap and not has_scaling
    # Whether ufunc on memmap return memmap
    viral_memmap = memmap_after_ufunc()
    with InTemporaryDirectory():
        with open(fname, 'wb') as fobj:
            fobj.write(b' ' * offset)
            fobj.write(arr.tostring(order='F'))
        for mmap, expected_mode in (
                # mmap value, expected memmap mode
                # mmap=None -> no mmap value
                # expected mode=None -> no memmap returned
                (None, 'c'),
                (True, 'c'),
                ('c', 'c'),
                ('r', 'r'),
                (False, None)):
            kwargs = {}
            if mmap is not None:
                kwargs['mmap'] = mmap
            prox = proxy_class(fname, hdr, **kwargs)
            unscaled = prox.get_unscaled()
            back_data = np.asanyarray(prox)
            unscaled_is_mmap = isinstance(unscaled, np.memmap)
            back_is_mmap =  isinstance(back_data, np.memmap)
            if expected_mode is None:
                assert_false(unscaled_is_mmap)
                assert_false(back_is_mmap)
            else:
                assert_equal(unscaled_is_mmap,
                             viral_memmap or unscaled_really_mmap)
                assert_equal(back_is_mmap,
                             viral_memmap or scaled_really_mmap)
                if scaled_really_mmap:
                    assert_equal(back_data.mode, expected_mode)
            del prox, back_data
            # Check that mmap is keyword-only
            assert_raises(TypeError, proxy_class, fname, hdr, True)
            # Check invalid values raise error
            assert_raises(ValueError, proxy_class, fname, hdr, mmap='rw')
            assert_raises(ValueError, proxy_class, fname, hdr, mmap='r+')


# An image opener class which counts how many instances of itself have been
# created
class CountingImageOpener(ImageOpener):
    num_openers = 0
    def __init__(self, *args, **kwargs):
        super(CountingImageOpener, self).__init__(*args, **kwargs)
        CountingImageOpener.num_openers += 1


def _count_ImageOpeners(proxy, data, voxels):
    CountingImageOpener.num_openers = 0
    # expected data is defined in the test_keep_file_open_* tests
    for i in range(voxels.shape[0]):
        x, y, z = [int(c) for c in voxels[i, :]]
        assert proxy[x, y, z] == x * 100 + y * 10 + z
    return CountingImageOpener.num_openers


@contextlib.contextmanager
def patch_keep_file_open_default(value):
    # Patch arrayproxy.KEEP_FILE_OPEN_DEFAULT with the given value
    with mock.patch('nibabel.arrayproxy.KEEP_FILE_OPEN_DEFAULT', value):
        yield


def test_keep_file_open_true_false_invalid():
    # Test the behaviour of the keep_file_open __init__ flag, when it is set to
    # True or False. Expected behaviour is as follows:
    # keep_open | igzip present    | persist ImageOpener | igzip.drop_handles
    #           | and is gzip file |                     |
    # ----------|------------------|---------------------|-------------------
    # False     | False            | False               | n/a
    # False     | True             | True                | True
    # True      | False            | True                | n/a
    # True      | True             | True                | False
    # 'auto'    | False            | False               | n/a
    # 'auto'    | True             | True                | False
    #
    # Each test tuple contains:
    #  - file type - gzipped ('gz') or not ('bin'), or an open file handle
    #    ('open')
    #  - keep_file_open value passed to ArrayProxy
    #  - whether or not indexed_gzip is present
    #  - expected value for internal ArrayProxy._persist_opener flag
    #  - expected value for internal ArrayProxy._keep_file_open flag
    tests = [
        # open file handle - kfo and have_igzip are both irrelevant
        ('open', False,  False, False, False),
        ('open', False,  True,  False, False),
        ('open', True,   False, False, False),
        ('open', True,   True,  False, False),
        ('open', 'auto', False, False, False),
        ('open', 'auto', True,  False, False),
        # non-gzip file - have_igzip is irrelevant, decision should be made
        # solely from kfo flag
        ('bin', False,  False, False, False),
        ('bin', False,  True,  False, False),
        ('bin', True,   False, True,  True),
        ('bin', True,   True,  True,  True),
        ('bin', 'auto', False, False, False),
        ('bin', 'auto', True,  False, False),
        # gzip file. If igzip is present, we persist the ImageOpener. If kfo
        # is 'auto':
        #  - if igzip is present, kfo -> True
        #  - otherwise, kfo -> False
        ('gz', False,  False, False, False),
        ('gz', False,  True,  True,  False),
        ('gz', True,   False, True,  True),
        ('gz', True,   True,  True,  True),
        ('gz', 'auto', False, False, False),
        ('gz', 'auto', True,  True,  True)]

    dtype = np.float32
    data  = np.arange(1000, dtype=dtype).reshape((10, 10, 10))
    voxels = np.random.randint(0, 10, (10, 3))

    for test in tests:
        filetype, kfo, have_igzip, exp_persist, exp_kfo = test
        with InTemporaryDirectory(), \
             mock.patch('nibabel.openers.ImageOpener', CountingImageOpener), \
             patch_indexed_gzip(have_igzip):
            fname = 'testdata.{}'.format(filetype)
            # create the test data file
            if filetype == 'gz':
                with gzip.open(fname, 'wb') as fobj:
                    fobj.write(data.tostring(order='F'))
            else:
                with open(fname, 'wb') as fobj:
                    fobj.write(data.tostring(order='F'))
            # pass in a file name or open file handle. If the latter, we open
            # two file handles, because we're going to create two proxies
            # below.
            if filetype == 'open':
                fobj1 = open(fname, 'rb')
                fobj2 = open(fname, 'rb')
            else:
                fobj1 = fname
                fobj2 = fname
            try:
                proxy = ArrayProxy(fobj1, ((10, 10, 10), dtype),
                                   keep_file_open=kfo)
                # We also test that we get the same behaviour when the
                # KEEP_FILE_OPEN_DEFAULT flag is changed
                with patch_keep_file_open_default(kfo):
                    proxy_def = ArrayProxy(fobj2, ((10, 10, 10), dtype))
                # check internal flags
                assert proxy._persist_opener == exp_persist
                assert proxy._keep_file_open == exp_kfo
                assert proxy_def._persist_opener == exp_persist
                assert proxy_def._keep_file_open == exp_kfo
                # check persist_opener behaviour - whether one imageopener is
                # created for the lifetime of the ArrayProxy, or one is
                # created on each access
                if exp_persist:
                    assert _count_ImageOpeners(proxy, data, voxels) == 1
                    assert _count_ImageOpeners(proxy_def, data, voxels) == 1
                else:
                    assert _count_ImageOpeners(proxy, data, voxels) == 10
                    assert _count_ImageOpeners(proxy_def, data, voxels) == 10
                # if indexed_gzip is active, check that the file object was
                # created correctly - the _opener.fobj will be a
                # MockIndexedGzipFile, defined in test_openers.py
                if filetype == 'gz' and have_igzip:
                    assert proxy._opener.fobj._drop_handles == (not exp_kfo)
                # if we were using an open file handle, check that the proxy
                # didn't close it
                if filetype == 'open':
                    assert not fobj1.closed
                    assert not fobj2.closed
            finally:
                del proxy
                del proxy_def
                if filetype == 'open':
                    fobj1.close()
                    fobj2.close()
    # Test invalid values of keep_file_open
    with InTemporaryDirectory():
        fname = 'testdata'
        with open(fname, 'wb') as fobj:
            fobj.write(data.tostring(order='F'))
        with assert_raises(ValueError):
            ArrayProxy(fname, ((10, 10, 10), dtype), keep_file_open=55)
        with assert_raises(ValueError):
            ArrayProxy(fname, ((10, 10, 10), dtype), keep_file_open='autob')
        with assert_raises(ValueError):
            ArrayProxy(fname, ((10, 10, 10), dtype), keep_file_open='cauto')


def test_pickle_lock():
    # Test that ArrayProxy can be pickled, and that thread lock is created

    def islock(l):
        # isinstance doesn't work on threading.Lock?
        return hasattr(l, 'acquire') and hasattr(l, 'release')

    proxy = ArrayProxy('dummyfile', ((10, 10, 10), np.float32))
    assert islock(proxy._lock)
    pickled = pickle.dumps(proxy)
    unpickled = pickle.loads(pickled)
    assert islock(unpickled._lock)
    assert proxy._lock is not unpickled._lock
