from __future__ import division, print_function, absolute_import

from tempfile import mkdtemp, mktemp
import os
import shutil

import numpy as np
from numpy import array, transpose, pi
from numpy.testing import (assert_equal,
                           assert_array_equal, assert_array_almost_equal)
import pytest
from pytest import raises as assert_raises

import scipy.sparse
from scipy.io.mmio import mminfo, mmread, mmwrite

parametrize_args = [('integer', 'int'),
                    ('unsigned-integer', 'uint')]


class TestMMIOArray(object):
    def setup_method(self):
        self.tmpdir = mkdtemp()
        self.fn = os.path.join(self.tmpdir, 'testfile.mtx')

    def teardown_method(self):
        shutil.rmtree(self.tmpdir)

    def check(self, a, info):
        mmwrite(self.fn, a)
        assert_equal(mminfo(self.fn), info)
        b = mmread(self.fn)
        assert_array_almost_equal(a, b)

    def check_exact(self, a, info):
        mmwrite(self.fn, a)
        assert_equal(mminfo(self.fn), info)
        b = mmread(self.fn)
        assert_equal(a, b)

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_integer(self, typeval, dtype):
        self.check_exact(array([[1, 2], [3, 4]], dtype=dtype),
                         (2, 2, 4, 'array', typeval, 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_32bit_integer(self, typeval, dtype):
        a = array([[2**31-1, 2**31-2], [2**31-3, 2**31-4]], dtype=dtype)
        self.check_exact(a, (2, 2, 4, 'array', typeval, 'general'))

    def test_64bit_integer(self):
        a = array([[2**31, 2**32], [2**63-2, 2**63-1]], dtype=np.int64)
        if (np.intp(0).itemsize < 8):
            assert_raises(OverflowError, mmwrite, self.fn, a)
        else:
            self.check_exact(a, (2, 2, 4, 'array', 'integer', 'general'))

    def test_64bit_unsigned_integer(self):
        a = array([[2**31, 2**32], [2**64-2, 2**64-1]], dtype=np.uint64)
        self.check_exact(a, (2, 2, 4, 'array', 'unsigned-integer', 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_upper_triangle_integer(self, typeval, dtype):
        self.check_exact(array([[0, 1], [0, 0]], dtype=dtype),
                         (2, 2, 4, 'array', typeval, 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_lower_triangle_integer(self, typeval, dtype):
        self.check_exact(array([[0, 0], [1, 0]], dtype=dtype),
                         (2, 2, 4, 'array', typeval, 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_rectangular_integer(self, typeval, dtype):
        self.check_exact(array([[1, 2, 3], [4, 5, 6]], dtype=dtype),
                         (2, 3, 6, 'array', typeval, 'general'))

    def test_simple_rectangular_float(self):
        self.check([[1, 2], [3.5, 4], [5, 6]],
                   (3, 2, 6, 'array', 'real', 'general'))

    def test_simple_float(self):
        self.check([[1, 2], [3, 4.0]],
                   (2, 2, 4, 'array', 'real', 'general'))

    def test_simple_complex(self):
        self.check([[1, 2], [3, 4j]],
                   (2, 2, 4, 'array', 'complex', 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_symmetric_integer(self, typeval, dtype):
        self.check_exact(array([[1, 2], [2, 4]], dtype=dtype),
                         (2, 2, 4, 'array', typeval, 'symmetric'))

    def test_simple_skew_symmetric_integer(self):
        self.check_exact([[0, 2], [-2, 0]],
                         (2, 2, 4, 'array', 'integer', 'skew-symmetric'))

    def test_simple_skew_symmetric_float(self):
        self.check(array([[0, 2], [-2.0, 0.0]], 'f'),
                   (2, 2, 4, 'array', 'real', 'skew-symmetric'))

    def test_simple_hermitian_complex(self):
        self.check([[1, 2+3j], [2-3j, 4]],
                   (2, 2, 4, 'array', 'complex', 'hermitian'))

    def test_random_symmetric_float(self):
        sz = (20, 20)
        a = np.random.random(sz)
        a = a + transpose(a)
        self.check(a, (20, 20, 400, 'array', 'real', 'symmetric'))

    def test_random_rectangular_float(self):
        sz = (20, 15)
        a = np.random.random(sz)
        self.check(a, (20, 15, 300, 'array', 'real', 'general'))


class TestMMIOSparseCSR(TestMMIOArray):
    def setup_method(self):
        self.tmpdir = mkdtemp()
        self.fn = os.path.join(self.tmpdir, 'testfile.mtx')

    def teardown_method(self):
        shutil.rmtree(self.tmpdir)

    def check(self, a, info):
        mmwrite(self.fn, a)
        assert_equal(mminfo(self.fn), info)
        b = mmread(self.fn)
        assert_array_almost_equal(a.todense(), b.todense())

    def check_exact(self, a, info):
        mmwrite(self.fn, a)
        assert_equal(mminfo(self.fn), info)
        b = mmread(self.fn)
        assert_equal(a.todense(), b.todense())

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_integer(self, typeval, dtype):
        self.check_exact(scipy.sparse.csr_matrix([[1, 2], [3, 4]], dtype=dtype),
                         (2, 2, 4, 'coordinate', typeval, 'general'))

    def test_32bit_integer(self):
        a = scipy.sparse.csr_matrix(array([[2**31-1, -2**31+2],
                                           [2**31-3, 2**31-4]],
                                          dtype=np.int32))
        self.check_exact(a, (2, 2, 4, 'coordinate', 'integer', 'general'))

    def test_64bit_integer(self):
        a = scipy.sparse.csr_matrix(array([[2**32+1, 2**32+1],
                                           [-2**63+2, 2**63-2]],
                                          dtype=np.int64))
        if (np.intp(0).itemsize < 8):
            assert_raises(OverflowError, mmwrite, self.fn, a)
        else:
            self.check_exact(a, (2, 2, 4, 'coordinate', 'integer', 'general'))

    def test_32bit_unsigned_integer(self):
        a = scipy.sparse.csr_matrix(array([[2**31-1, 2**31-2],
                                           [2**31-3, 2**31-4]],
                                          dtype=np.uint32))
        self.check_exact(a, (2, 2, 4, 'coordinate', 'unsigned-integer', 'general'))

    def test_64bit_unsigned_integer(self):
        a = scipy.sparse.csr_matrix(array([[2**32+1, 2**32+1],
                                           [2**64-2, 2**64-1]],
                                          dtype=np.uint64))
        self.check_exact(a, (2, 2, 4, 'coordinate', 'unsigned-integer', 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_upper_triangle_integer(self, typeval, dtype):
        self.check_exact(scipy.sparse.csr_matrix([[0, 1], [0, 0]], dtype=dtype),
                         (2, 2, 1, 'coordinate', typeval, 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_lower_triangle_integer(self, typeval, dtype):
        self.check_exact(scipy.sparse.csr_matrix([[0, 0], [1, 0]], dtype=dtype),
                         (2, 2, 1, 'coordinate', typeval, 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_rectangular_integer(self, typeval, dtype):
        self.check_exact(scipy.sparse.csr_matrix([[1, 2, 3], [4, 5, 6]], dtype=dtype),
                         (2, 3, 6, 'coordinate', typeval, 'general'))

    def test_simple_rectangular_float(self):
        self.check(scipy.sparse.csr_matrix([[1, 2], [3.5, 4], [5, 6]]),
                   (3, 2, 6, 'coordinate', 'real', 'general'))

    def test_simple_float(self):
        self.check(scipy.sparse.csr_matrix([[1, 2], [3, 4.0]]),
                   (2, 2, 4, 'coordinate', 'real', 'general'))

    def test_simple_complex(self):
        self.check(scipy.sparse.csr_matrix([[1, 2], [3, 4j]]),
                   (2, 2, 4, 'coordinate', 'complex', 'general'))

    @pytest.mark.parametrize('typeval, dtype', parametrize_args)
    def test_simple_symmetric_integer(self, typeval, dtype):
        self.check_exact(scipy.sparse.csr_matrix([[1, 2], [2, 4]], dtype=dtype),
                         (2, 2, 3, 'coordinate', typeval, 'symmetric'))

    def test_simple_skew_symmetric_integer(self):
        self.check_exact(scipy.sparse.csr_matrix([[1, 2], [-2, 4]]),
                         (2, 2, 3, 'coordinate', 'integer', 'skew-symmetric'))

    def test_simple_skew_symmetric_float(self):
        self.check(scipy.sparse.csr_matrix(array([[1, 2], [-2.0, 4]], 'f')),
                   (2, 2, 3, 'coordinate', 'real', 'skew-symmetric'))

    def test_simple_hermitian_complex(self):
        self.check(scipy.sparse.csr_matrix([[1, 2+3j], [2-3j, 4]]),
                   (2, 2, 3, 'coordinate', 'complex', 'hermitian'))

    def test_random_symmetric_float(self):
        sz = (20, 20)
        a = np.random.random(sz)
        a = a + transpose(a)
        a = scipy.sparse.csr_matrix(a)
        self.check(a, (20, 20, 210, 'coordinate', 'real', 'symmetric'))

    def test_random_rectangular_float(self):
        sz = (20, 15)
        a = np.random.random(sz)
        a = scipy.sparse.csr_matrix(a)
        self.check(a, (20, 15, 300, 'coordinate', 'real', 'general'))

    def test_simple_pattern(self):
        a = scipy.sparse.csr_matrix([[0, 1.5], [3.0, 2.5]])
        p = np.zeros_like(a.todense())
        p[a.todense() > 0] = 1
        info = (2, 2, 3, 'coordinate', 'pattern', 'general')
        mmwrite(self.fn, a, field='pattern')
        assert_equal(mminfo(self.fn), info)
        b = mmread(self.fn)
        assert_array_almost_equal(p, b.todense())


_32bit_integer_dense_example = '''\
%%MatrixMarket matrix array integer general
2  2
2147483647
2147483646
2147483647
2147483646
'''

_32bit_integer_sparse_example = '''\
%%MatrixMarket matrix coordinate integer symmetric
2  2  2
1  1  2147483647
2  2  2147483646
'''

_64bit_integer_dense_example = '''\
%%MatrixMarket matrix array integer general
2  2
          2147483648
-9223372036854775806
         -2147483648
 9223372036854775807
'''

_64bit_integer_sparse_general_example = '''\
%%MatrixMarket matrix coordinate integer general
2  2  3
1  1           2147483648
1  2  9223372036854775807
2  2  9223372036854775807
'''

_64bit_integer_sparse_symmetric_example = '''\
%%MatrixMarket matrix coordinate integer symmetric
2  2  3
1  1            2147483648
1  2  -9223372036854775807
2  2   9223372036854775807
'''

_64bit_integer_sparse_skew_example = '''\
%%MatrixMarket matrix coordinate integer skew-symmetric
2  2  3
1  1            2147483648
1  2  -9223372036854775807
2  2   9223372036854775807
'''

_over64bit_integer_dense_example = '''\
%%MatrixMarket matrix array integer general
2  2
         2147483648
9223372036854775807
         2147483648
9223372036854775808
'''

_over64bit_integer_sparse_example = '''\
%%MatrixMarket matrix coordinate integer symmetric
2  2  2
1  1            2147483648
2  2  19223372036854775808
'''

class TestMMIOReadLargeIntegers(object):
    def setup_method(self):
        self.tmpdir = mkdtemp()
        self.fn = os.path.join(self.tmpdir, 'testfile.mtx')

    def teardown_method(self):
        shutil.rmtree(self.tmpdir)

    def check_read(self, example, a, info, dense, over32, over64):
        with open(self.fn, 'w') as f:
            f.write(example)
        assert_equal(mminfo(self.fn), info)
        if (over32 and (np.intp(0).itemsize < 8)) or over64:
            assert_raises(OverflowError, mmread, self.fn)
        else:
            b = mmread(self.fn)
            if not dense:
                b = b.todense()
            assert_equal(a, b)

    def test_read_32bit_integer_dense(self):
        a = array([[2**31-1, 2**31-1],
                   [2**31-2, 2**31-2]], dtype=np.int64)
        self.check_read(_32bit_integer_dense_example,
                        a,
                        (2, 2, 4, 'array', 'integer', 'general'),
                        dense=True,
                        over32=False,
                        over64=False)

    def test_read_32bit_integer_sparse(self):
        a = array([[2**31-1, 0],
                   [0, 2**31-2]], dtype=np.int64)
        self.check_read(_32bit_integer_sparse_example,
                        a,
                        (2, 2, 2, 'coordinate', 'integer', 'symmetric'),
                        dense=False,
                        over32=False,
                        over64=False)

    def test_read_64bit_integer_dense(self):
        a = array([[2**31, -2**31],
                   [-2**63+2, 2**63-1]], dtype=np.int64)