from fits import *
import numpy as np
import unittest


class TestSlices(unittest.TestCase):
	def setUp(self):
		T = tabledata()
		T.x1 = np.arange(10).astype(int)
		T.x2 = np.arange(10).astype(float)
		T.x3 = [i for i in range(10)]
		T.x4 = tuple([i for i in range(10)])
		self.assertEqual(len(T), 10)
		T.about()
		self.T = T

	def testCut(self):
		T1 = self.T.copy()
		T1.about()

		T1.cut(np.arange(10) < 4)
		T1.about()
		self.assertEqual(len(T1), 4)

	def testIter(self):
		for i in range(len(self.T)):
			self.T[i].about()
		for t in self.T:
			t.about()

	def testNumpySlicing(self):
		A = np.arange(10)
		print A[:3]
		print A[3:6]
		print A[7:]
		print A[1:7:2]

		A = list(A)
		print A[:3]
		print A[3:6]
		print A[7:]
		print A[1:7:2]

		A = np.arange(10)
		A = A[:,np.newaxis].repeat(2, axis=1)
		print A[:3]
		print A[3:6]
		print A[7:]
		print A[1:7:2]

	def testSlice2(self):
		T1 = self.T[self.T.x1 < 4]
		self.assertEqual(len(T1), 4)
		T1.about()
		self.assertTrue(all(T1.x1 == np.array([0,1,2,3])))
		self.assertEqual(T1.x3, [0,1,2,3])

		T1 = self.T[np.flatnonzero(self.T.x1 < 4)]
		self.assertEqual(len(T1), 4)
		T1.about()
		self.assertTrue(all(T1.x1 == np.array([0,1,2,3])))
		self.assertEqual(T1.x3, [0,1,2,3])

	def testSlice1(self):
		T1 = self.T[:3]
		self.assertEqual(len(T1), 3)
		T1.about()
		self.assertTrue(all(T1.x1 == np.arange(3)))
		self.assertTrue(all(T1.x2 == np.arange(3).astype(float)))
		self.assertEqual(T1.x3, [0,1,2])
		self.assertEqual(T1.x4, (0,1,2))

		T2 = self.T[3:6]
		self.assertEqual(len(T2), 3)
		T2.about()
		self.assertEqual(T2.x3, [3,4,5])
		self.assertEqual(T2.x4, (3,4,5))

		T3 = self.T[7:]
		self.assertEqual(len(T3), 3)
		T3.about()
		self.assertEqual(T3.x3, [7,8,9])

		T4 = self.T[1:7:2]
		self.assertEqual(len(T4), 3)
		T4.about()
		self.assertEqual(T4.x3, [1,3,5])


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



