from secml.array.tests import CArrayTestCases

from secml.utils import fm
from secml.array import CArray


class TestCArraySaveLoad(CArrayTestCases):
    """Unit test for CArray SAVE/LOAD methods.

    The test cases have been splitted in multiple methods to make unittests
    call tearDown after each one, resetting test files even after errors.

    """
    def setUp(self):

        super(TestCArraySaveLoad, self).setUp()

        self.test_file = fm.join(fm.abspath(__file__), 'test.txt')
        self.test_file_2 = fm.join(fm.abspath(__file__), 'test2.txt')

    def test_save_load_sparse(self):
        """Test save/load of CArray"""
        self.logger.info(
            "UNITTEST - CArray - Testing save/load for sparse matrix")

        self.array_sparse.save(self.test_file)

        # Saving to a file handle is not supported for sparse arrays
        with self.assertRaises(NotImplementedError):
            with open(self.test_file_2, 'w') as f:
                self.array_sparse.save(f)

        loaded_array_sparse = CArray.load(
            self.test_file, arrayformat='sparse', dtype=int)

        self.assertFalse((loaded_array_sparse != self.array_sparse).any(),
                         "Saved and loaded arrays (sparse) are not equal!")

    def test_save_load_sparse_conversion(self):
        """Test save/load of CArray"""
        # Array should be stored and loaded correctly whatever sparse format
        self.array_sparse._data._data = self.array_sparse._data.todok()

        self.array_sparse.save(self.test_file)

        # Saving to a file handle is not supported for sparse arrays
        with self.assertRaises(NotImplementedError):
            with open(self.test_file_2, 'w') as f:
                self.array_sparse.save(f)

        loaded_array_sparse = CArray.load(
            self.test_file, arrayformat='sparse', dtype=int)

        self.assertFalse((loaded_array_sparse != self.array_sparse).any(),
                         "Saved and loaded arrays (sparse) are not equal!")

    def test_save_load_dense(self):
        """Test save/load of CArray"""
        self.logger.info(
            "UNITTEST - CSparse - Testing save/load for dense matrix")

        self.array_dense.save(self.test_file, overwrite=True)

        loaded_array_dense = CArray.load(
            self.test_file, arrayformat='dense', dtype=int)

        self.assertFalse((loaded_array_dense != self.array_dense).any(),
                         "Saved and loaded arrays (sparse) are not equal!")

        # Only 'dense' and 'sparse' arrayformat are supported
        with self.assertRaises(ValueError):
            CArray.load(self.test_file, arrayformat='test')

    def tearDown(self):
        # Remove test file(s) if exist
        try:
            fm.remove_file(self.test_file)
            fm.remove_file(self.test_file_2)
        except (OSError, IOError) as e:
            if e.errno != 2:
                raise e


if __name__ == '__main__':
    CArrayTestCases.main()
