import unittest
from petsc4py import PETSc

# --------------------------------------------------------------------


def allclose(seq1, seq2):
    for v1, v2 in zip(seq1, seq2):
        if abs(v1 - v2) > 1e-5:
            return False
    return True


class TestNullSpace(unittest.TestCase):
    def setUp(self):
        u1 = PETSc.Vec().createSeq(3)
        u2 = PETSc.Vec().createSeq(3)
        u1[0], u1[1], u1[2] = [1, 2, 0]
        u1.normalize()
        u2[0], u2[1], u2[2] = [2, -1, 0]
        u2.normalize()
        basis = [u1, u2]
        nullsp = PETSc.NullSpace().create(False, basis, comm=PETSc.COMM_SELF)
        self.basis = basis
        self.nullsp = nullsp

    def tearDown(self):
        self.basis = None
        self.nullsp = None
        PETSc.garbage_cleanup()

    def _remove(self):
        v = PETSc.Vec().createSeq(3)
        v[0], v[1], v[2] = [7, 8, 9]
        w = v.copy()
        self.nullsp.remove(w)
        return (v, w)

    def testRemove(self):
        v, w = self._remove()
        self.assertTrue(allclose(v.array, [7, 8, 9]))
        self.assertTrue(allclose(w.array, [0, 0, 9]))
        del v, w

    def testRemoveInplace(self):
        v, w = self._remove()
        self.nullsp.remove(v)
        self.assertTrue(v.equal(w))
        del v, w

    def testRemoveWithFunction(self):
        def myremove(nsp, vec):
            vec.setArray([1, 2, 3])

        self.nullsp.setFunction(myremove)
        v, w = self._remove()
        self.assertTrue(allclose(v.array, [7, 8, 9]))
        self.assertTrue(allclose(w.array, [1, 2, 3]))
        self.nullsp.remove(v)
        self.assertTrue(allclose(v.array, [1, 2, 3]))
        self.nullsp.setFunction(None)
        self.testRemove()

    def testGetSetFunction(self):
        def rem(nsp, vec):
            vec.set(0)

        self.nullsp.setFunction(rem)
        dct = self.nullsp.getDict()
        self.assertTrue(dct is not None)
        fun, a, kw = dct['__function__']
        self.assertTrue(fun is rem)
        self.nullsp.setFunction(None)
        fun = dct.get('__function__')
        self.assertTrue(fun is None)


# --------------------------------------------------------------------

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

# --------------------------------------------------------------------