""" Test imageio core functionality.
"""

import os
from zipfile import ZipFile

import numpy as np

from pytest import raises
from imageio.testing import run_tests_if_main, get_test_dir, need_internet

import imageio
from imageio import core
from imageio.core import get_remote_file


test_dir = get_test_dir()


_prepared = None


def _prepare():
    """ Create two dirs, one with one dataset and one with two datasets
    """
    need_internet()
    
    global _prepared
    if _prepared and os.path.isfile(_prepared[2]):
        return _prepared
    # Prepare sources
    fname1 = get_remote_file('images/dicom_sample1.zip')
    fname2 = get_remote_file('images/dicom_sample2.zip')
    dname1 = os.path.join(test_dir, 'dicom_sample1')
    dname2 = os.path.join(test_dir, 'dicom_sample2')
    # Extract zipfiles
    z = ZipFile(fname1)
    z.extractall(dname1)
    z.extractall(dname2)
    z = ZipFile(fname2)
    z.extractall(dname2)
    # Get arbitrary file names
    fname1 = os.path.join(dname1, os.listdir(dname1)[0])
    fname2 = os.path.join(dname2, os.listdir(dname2)[0])
    # Cache and return
    _prepared = dname1, dname2, fname1, fname2
    return dname1, dname2, fname1, fname2


def test_read_empty_dir():
    # Create an empty dir
    empty = os.path.join(test_dir, 'empty_dir')
    if not os.path.isdir(empty):
        os.mkdir(empty)
    
    # Tes that no format is found, but no error is raised
    request = core.Request(empty, 'ri')
    assert imageio.formats.search_read_format(request) is None
    

def test_selection():
    
    dname1, dname2, fname1, fname2 = _prepare()
    
    # Test that DICOM can examine file
    F = imageio.formats.search_read_format(core.Request(fname1, 'ri'))
    assert F.name == 'DICOM'
    assert F is imageio.formats['DICOM']
    
    # Test that we cannot save
    request = core.Request(os.path.join(test_dir, 'test.dcm'), 'wi')
    assert not F.can_write(request)
    
    # Test fail on wrong file
    fname2 = fname1 + '.fake'
    bb = open(fname1, 'rb').read()
    bb = bb[:128] + b'XXXX' + bb[132:]
    open(fname2, 'wb').write(bb)
    raises(Exception, F.get_reader, core.Request(fname2, 'ri'))
    
    # Test special files with other formats
    im = imageio.imread(get_remote_file('images/dicom_file01.dcm'))
    assert im.shape == (512, 512)
    im = imageio.imread(get_remote_file('images/dicom_file03.dcm'))
    assert im.shape == (512, 512)
    im = imageio.imread(get_remote_file('images/dicom_file04.dcm'))
    assert im.shape == (512, 512)
    
    # Expected fails
    fname = get_remote_file('images/dicom_file90.dcm')
    raises(RuntimeError, imageio.imread, fname)  # 1.2.840.10008.1.2.4.91 
    fname = get_remote_file('images/dicom_file91.dcm')
    raises(RuntimeError, imageio.imread, fname)  # not pixel data
    
    # This one *should* work, but does not, see issue #18
    try:
        imageio.imread(get_remote_file('images/dicom_file02.dcm'))
    except Exception:
        pass


def test_progress():
    
    dname1, dname2, fname1, fname2 = _prepare()
    
    imageio.imread(fname1, progress=True)
    imageio.imread(fname1, progress=core.StdoutProgressIndicator('test'))
    imageio.imread(fname1, progress=None)
    raises(ValueError, imageio.imread, fname1, progress=3)


def test_different_read_modes():
    
    dname1, dname2, fname1, fname2 = _prepare()
    
    for fname, dname, n in [(fname1, dname1, 1), (fname2, dname2, 2)]:
        
        # Test imread()
        im = imageio.imread(fname)
        assert isinstance(im, np.ndarray)
        assert im.shape == (512, 512)
        
        # Test mimread()
        ims = imageio.mimread(fname)
        assert isinstance(ims, list)
        assert ims[0].shape == im.shape
        assert len(ims) > 1
        #    
        ims2 = imageio.mimread(dname)
        assert len(ims) == len(ims2)
        
        # Test volread()
        vol = imageio.volread(dname)
        assert vol.ndim == 3
        assert vol.shape[0] > 10
        assert vol.shape[1:] == (512, 512)
        #
        vol2 = imageio.volread(fname)  # fname works as well
        assert (vol == vol2).all()
        
        # Test mvolread()
        vols = imageio.mvolread(dname)
        assert isinstance(vols, list)
        assert len(vols) == n
        assert vols[0].shape == vol.shape
        assert sum([v.shape[0] for v in vols]) == len(ims)
    

def test_different_read_modes_with_readers():
    
    dname1, dname2, fname1, fname2 = _prepare()
    
    for fname, dname, n in [(fname1, dname1, 1), (fname2, dname2, 2)]:
        
        # Test imread()
        R = imageio.read(fname, 'DICOM', 'i')
        assert len(R) == 1
        assert isinstance(R.get_meta_data(), dict)
        assert isinstance(R.get_meta_data(0), dict)
        
        # Test mimread()
        R = imageio.read(fname, 'DICOM', 'I')
        if n == 1:
            assert len(R) > 10
        else:
            assert len(R) == 20 + 25
        assert isinstance(R.get_meta_data(), dict)
        assert isinstance(R.get_meta_data(0), dict)
        
        # Test volread()
        R = imageio.read(fname, 'DICOM', 'v')
        assert len(R) == n  # we ask for one, but get an honest number
        assert isinstance(R.get_meta_data(), dict)
        assert isinstance(R.get_meta_data(0), dict)
        
        # Test mvolread()
        R = imageio.read(fname, 'DICOM', 'V')
        assert len(R) == n 
        assert isinstance(R.get_meta_data(), dict)
        assert isinstance(R.get_meta_data(0), dict)
        
        # Touch DicomSeries objects
        assert repr(R._series[0])
        assert R._series[0].description
        assert len(R._series[0].sampling) == 3
        
        R = imageio.read(fname, 'DICOM', '?')
        raises(RuntimeError, R.get_length)


run_tests_if_main()
