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

import os
import gc
import shutil

import numpy as np

import imageio
from imageio.core import Format, FormatManager, Request
from imageio.core import get_remote_file


test_dir = get_test_dir()


def setup_module():
    imageio.formats.sort()


def teardown_module():
    imageio.formats.sort()


class MyFormat(Format):
    """ TEST DOCS """
    _closed = []
    
    def _can_read(self, request):
        return request.filename.lower().endswith(self.extensions + ('.haha', ))
    
    def _can_write(self, request):
        return request.filename.lower().endswith(self.extensions + ('.haha', ))
    
    class Reader(Format.Reader):
        _failmode = False
        _stream_mode = False
        
        def _open(self):
            self._read_frames = 0
        
        def _close(self):
            self.format._closed.append(id(self))
        
        def _get_length(self):
            if self._stream_mode:
                return np.inf
            return 3
        
        def _get_data(self, index):
            if self._failmode == 2:
                raise IndexError()
            elif self._failmode:
                return 'not an array', {}
            elif self._stream_mode and self._read_frames >= 5:
                raise IndexError()  # Mark end of stream
            else:
                self._read_frames += 1
                return np.ones((10, 10)) * index, self._get_meta_data(index)
        
        def _get_meta_data(self, index):
            if self._failmode:
                return 'not a dict'
            return {'index': index}
            
    class Writer(Format.Writer):
        
        def _open(self):
            self._written_data = []
            self._written_meta = []
            self._meta = None
        
        def _close(self):
            self.format._closed.append(id(self))
        
        def _append_data(self, im, meta):
            self._written_data.append(im)
            self._written_meta.append(meta)
        
        def _set_meta_data(self, meta):
            self._meta = meta


def test_format():
    """ Test the working of the Format class """
    
    filename1 = get_remote_file('images/chelsea.png', test_dir)
    filename2 = filename1 + '.out'
    
    # Test basic format creation
    F = Format('testname', 'test description', 'foo bar spam')
    assert F.name == 'TESTNAME'
    assert F.description == 'test description'
    assert F.name in repr(F)
    assert F.name in F.doc
    assert str(F) == F.doc
    assert set(F.extensions) == set(['.foo', '.bar', '.spam'])
    
    # Test setting extensions
    F1 = Format('test', '', 'foo bar spam')
    F2 = Format('test', '', 'foo, bar,spam')
    F3 = Format('test', '', ['foo', 'bar', 'spam'])
    F4 = Format('test', '', '.foo .bar .spam')
    for F in (F1, F2, F3, F4):
        assert set(F.extensions) == set(['.foo', '.bar', '.spam'])
    # Fail
    raises(ValueError, Format, 'test', '', 3)  # not valid ext
    raises(ValueError, Format, 'test', '', '', 3)  # not valid mode
    raises(ValueError, Format, 'test', '', '', 'x')  # not valid mode
    
    # Test subclassing
    F = MyFormat('test', '', modes='i')
    assert 'TEST DOCS' in F.doc
    
    # Get and check reader and write classes
    R = F.get_reader(Request(filename1, 'ri'))
    W = F.get_writer(Request(filename2, 'wi'))
    assert isinstance(R, MyFormat.Reader)
    assert isinstance(W, MyFormat.Writer)
    assert R.format is F
    assert W.format is F
    assert R.request.filename == filename1
    assert W.request.filename == filename2
    # Fail
    raises(RuntimeError, F.get_reader, Request(filename1, 'rI'))
    raises(RuntimeError, F.get_writer, Request(filename2, 'wI'))
    
    # Use as context manager
    with R:
        pass
    with W:
        pass
    # Objects are now closed, cannot be used
    assert R.closed
    assert W.closed
    #
    raises(RuntimeError, R.__enter__)
    raises(RuntimeError, W.__enter__)
    #
    raises(RuntimeError, R.get_data, 0)
    raises(RuntimeError, W.append_data, np.zeros((10, 10)))
    
    # Test __del__
    R = F.get_reader(Request(filename1, 'ri'))
    W = F.get_writer(Request(filename2, 'wi'))
    ids = id(R), id(W)
    F._closed[:] = []
    del R
    del W
    gc.collect()  # Invoke __del__
    assert set(ids) == set(F._closed)


def test_reader_and_writer():
    
    # Prepare
    filename1 = get_remote_file('images/chelsea.png', test_dir)
    filename2 = filename1 + '.out'
    F = MyFormat('test', '', modes='i')
    
    # Test using reader
    n = 3
    R = F.get_reader(Request(filename1, 'ri'))
    assert len(R) == n
    ims = [im for im in R]
    assert len(ims) == n
    for i in range(3):
        assert ims[i][0, 0] == i
        assert ims[i].meta['index'] == i
    for i in range(3):
        assert R.get_meta_data(i)['index'] == i
    # Read next
    assert R.get_data(0)[0, 0] == 0
    assert R.get_next_data()[0, 0] == 1
    assert R.get_next_data()[0, 0] == 2
    # Fail
    R._failmode = 1
    raises(ValueError, R.get_data, 0)
    raises(ValueError, R.get_meta_data, 0)
    R._failmode = 2
    with raises(IndexError):
        [im for im in R]
    
    # Test writer no format
    raises(ValueError, imageio.get_writer, 'foo.unknownext')
    
    # Test streaming reader
    R = F.get_reader(Request(filename1, 'ri'))
    R._stream_mode = True
    assert R.get_length() == np.inf
    ims = [im for im in R]
    assert len(ims) == 5
    
    # Test using writer
    im1 = np.zeros((10, 10))
    im2 = imageio.core.Image(im1, {'foo': 1})
    W = F.get_writer(Request(filename2, 'wi'))
    W.append_data(im1)
    W.append_data(im2)
    W.append_data(im1, {'bar': 1})
    W.append_data(im2, {'bar': 1})
    # Test that no data is copies (but may be different views)
    assert len(W._written_data) == 4 
    for im in W._written_data:
        assert (im == im1).all()
    im1[2, 2] == 99
    for im in W._written_data:
        assert (im == im1).all()
    # Test meta
    assert W._written_meta[0] == {}
    assert W._written_meta[1] == {'foo': 1}
    assert W._written_meta[2] == {'bar': 1}
    assert W._written_meta[3] == {'foo': 1, 'bar': 1}
    #
    W.set_meta_data({'spam': 1})
    assert W._meta == {'spam': 1}
    # Fail
    raises(ValueError, W.append_data, 'not an array')
    raises(ValueError, W.append_data, im, 'not a dict')
    raises(ValueError, W.set_meta_data, 'not a dict')


def test_default_can_read_and_can_write():
    
    F = imageio.plugins.example.DummyFormat('test', '', 'foo bar', 'v')
    
    # Prepare files
    filename1 = os.path.join(test_dir, 'test')
    open(filename1 + '.foo', 'wb')
    open(filename1 + '.bar', 'wb')
    open(filename1 + '.spam', 'wb')
    
    # Test _can_read()
    assert F.can_read(Request(filename1 + '.foo', 'rv'))
    assert F.can_read(Request(filename1 + '.bar', 'r?'))
    assert not F.can_read(Request(filename1 + '.spam', 'r?'))
    assert not F.can_read(Request(filename1 + '.foo', 'ri'))
    
    # Test _can_write()
    assert F.can_write(Request(filename1 + '.foo', 'wv'))
    assert F.can_write(Request(filename1 + '.bar', 'w?'))
    assert not F.can_write(Request(filename1 + '.spam', 'w?'))
    assert not F.can_write(Request(filename1 + '.foo', 'wi'))


def test_format_selection():
    
    formats = imageio.formats
    fname1 = get_remote_file('images/chelsea.png', test_dir)
    fname2 = os.path.join(test_dir, 'test.selectext1')
    fname3 = os.path.join(test_dir, 'test.haha')
    open(fname2, 'wb')
    open(fname3, 'wb')
    
    # Test searchinhg for read / write format
    F = formats.search_read_format(Request(fname1, 'ri'))
    assert F is formats['PNG']
    F = formats.search_write_format(Request(fname1, 'wi'))
    assert F is formats['PNG']
    
    # Now with custom format
    format = MyFormat('test_selection', 'xx', 'selectext1', 'i')
    formats.add_format(format)
    
    # Select this format for files it said it could handle in extensions
    assert '.selectext1' in fname2
    F = formats.search_read_format(Request(fname2, 'ri'))
    assert F is format
    F = formats.search_write_format(Request(fname2, 'ri'))
    assert F is format
    
    # But this custom format also can deal with .haha files
    assert '.haha' in fname3
    F = formats.search_read_format(Request(fname3, 'ri'))
    assert F is format
    F = formats.search_write_format(Request(fname3, 'ri'))
    assert F is format


## Format manager

def test_format_manager():
    """ Test working of the format manager """
    
    formats = imageio.formats
    
    # Test basics of FormatManager
    assert isinstance(formats, FormatManager)
    assert len(formats) > 0
    assert 'FormatManager' in repr(formats)
    
    # Get docs
    smalldocs = str(formats)
    #fulldocs = formats.create_docs_for_all_formats()
    
    # Check each format ...
    for format in formats:
        #  That each format is indeed a Format
        assert isinstance(format, Format)
        # That they are mentioned
        assert format.name in smalldocs
        #assert format.name in fulldocs
    
    fname = get_remote_file('images/chelsea.png', test_dir)
    fname2 = fname[:-3] + 'noext'
    shutil.copy(fname, fname2)
    
    # Check getting
    F1 = formats['PNG']
    F2 = formats['.png']
    F3 = formats[fname2]  # will look in file itself
    assert F1 is F2
    assert F1 is F3
    # Check getting
    F1 = formats['DICOM']
    F2 = formats['.dcm']
    F3 = formats['dcm']  # If omitting dot, format is smart enough to try with
    assert F1 is F2
    assert F1 is F3
    # Fail
    raises(ValueError, formats.__getitem__, 678)  # must be str
    raises(IndexError, formats.__getitem__, '.nonexistentformat')
    
    # Adding a format
    myformat = Format('test', 'test description', 'testext1 testext2')
    formats.add_format(myformat)
    assert myformat in [f for f in formats]
    assert formats['testext1'] is myformat
    assert formats['testext2'] is myformat
    # Fail
    raises(ValueError, formats.add_format, 678)  # must be Format
    raises(ValueError, formats.add_format, myformat)  # cannot add twice
    
    # Adding a format with the same name
    myformat2 = Format('test', 'other description', 'foo bar')
    raises(ValueError, formats.add_format, myformat2)  # same name
    formats.add_format(myformat2, True)  # overwrite
    assert formats['test'] is not myformat
    assert formats['test'] is myformat2
    
    # Test show (we assume it shows correctly)
    formats.show()
    
#   # Potential
#   bytes = b'x' * 300
#   F = formats.search_read_format(Request(bytes, 'r?', dummy_potential=1))
#   assert F is formats['DUMMY']


def test_sorting_errors():
    
    with raises(TypeError):
        imageio.formats.sort(3)
    with raises(ValueError):
        imageio.formats.sort('foo,bar')
    with raises(ValueError):
        imageio.formats.sort('foo.png')


def test_default_order():
    
    assert imageio.formats['.tiff'].name == 'TIFF'
    assert imageio.formats['.png'].name == 'PNG-PIL'
    assert imageio.formats['.pfm'].name == 'PFM-FI'


def test_preferring_fi():
    
    # Prefer FI all the way
    imageio.formats.sort('-FI')
    
    assert imageio.formats['.tiff'].name == 'TIFF-FI'
    assert imageio.formats['.png'].name == 'PNG-FI'
    assert imageio.formats['.pfm'].name == 'PFM-FI'
    
    # This would be better
    imageio.formats.sort('TIFF', '-FI')
    assert imageio.formats['.tiff'].name == 'TIFF'
    assert imageio.formats['.png'].name == 'PNG-FI'
    assert imageio.formats['.pfm'].name == 'PFM-FI'


def test_preferring_arbitrary():
    
    # Normally, these exotic formats are somewhere in the back
    imageio.formats.sort()
    names = [f.name for f in imageio.formats]
    assert 'DICOM' not in names[:10]
    assert 'FFMPEG' not in names[:10]
    assert 'NPZ' not in names[:10]
    
    # But we can move them forward
    imageio.formats.sort('DICOM', 'FFMPEG', 'NPZ')
    names = [f.name for f in imageio.formats]
    assert names[0] == 'DICOM'
    assert names[1] == 'FFMPEG'
    assert names[2] == 'NPZ'
    
    # And back to normal ..
    imageio.formats.sort()
    names = [f.name for f in imageio.formats]
    assert 'DICOM' not in names[:10]
    assert 'FFMPEG' not in names[:10]
    assert 'NPZ' not in names[:10]


run_tests_if_main()
