"""
Test of neo.io.asciiimageio
"""
import os
import unittest
import quantities as pq
from neo.io import AsciiImageIO
import numpy as np


class TestAsciiImageIO(unittest.TestCase):

    def test_read_txt(self):
        img = ''
        img_list = []
        for frame in range(20):
            img_list.append([])
            for y in range(50):
                img_list[frame].append([])
                for x in range(50):
                    img += str(x)
                    img += '\t'
                    img_list[frame][y].append(x)
        img_list = np.array(img_list)
        file_name = "txt_test_file.txt"
        file = open(file_name, mode="w")
        file.write(str(img))
        file.close()

        object = AsciiImageIO(file_name='txt_test_file.txt',
                              nb_frame=20, nb_row=50, nb_column=50, units='V',
                              sampling_rate=1 * pq.Hz, spatial_scale=1 * pq.micrometer)
        block = object.read_block()
        self.assertEqual(len(block.segments), 1)
        self.assertEqual(len(block.segments[0].imagesequences), 1)
        self.assertEqual(block.segments[0].imagesequences[0].shape, (20, 50, 50))
        self.assertEqual(block.segments[0].imagesequences[0].any(), img_list.any())

        file.close()
        os.remove(file_name)


if __name__ == "__main__":
    unittest.main()