import unittest
import numpy

from numpy.testing import *
set_package_path()
from transformation import *
restore_path()

class TestDeformationCreation(NumpyTestCase):
  def test_1DCreation(self):
    size = (10,)
    points = numpy.array(((2,), (8, )))
    displacements = [0, (-4,)]

    field = denseDeformationFieldFromSparse(size, points, displacements)

    for point,displacement in zip(points, displacements):
      assert_almost_equal(field[tuple(point.tolist())], displacement)

  def test_2DCreation(self):
    size = (11, 11)
    points = numpy.array(((5., 6.), (5., 4.), (6., 5.), (4., 5.)))
    displacements = [(0, 1), (0, 1) , (0, -1), (0, -1)]

    field = denseDeformationFieldFromSparse(size, points, displacements)
    for point,displacement in zip(points, displacements):
      assert_almost_equal(field[tuple(point.tolist())], displacement)

  def test_2DCreation_bis(self):
    size = (10, 10)
    points = numpy.array(((5., 5.), (2., 8.), (8., 2.)))
    displacements = [0, (5, -4), (0, 2)]

    field = denseDeformationFieldFromSparse(size, points, displacements)
    for point,displacement in zip(points, displacements):
      assert_almost_equal(field[tuple(point.tolist())], displacement)

  def test_3DCreation(self):
    size = (10, 10, 10)
    points = numpy.array(((5., 5., 5.), (2., 8., 2.), (8., 2., 8.), (2., 2., 8.)))
    displacements = [0, (5, -4, 3), (0, 2, 0), -2]

    field = denseDeformationFieldFromSparse(size, points, displacements)

    for point,displacement in zip(points, displacements):
      assert_almost_equal(field[tuple(point.tolist())], displacement)

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