import sys
import os
import unittest
import numpy as np

import uLib

def vector4f0(v, target):
    diff = np.array(v) - np.array(target)
    diff[3] = 0 # ignoring w
    return np.all(np.abs(diff) < 0.001)


class TestMathMatrix(unittest.TestCase):

    def test_matrix(self):
        def check_1234(m2f):
            self.assertEqual(m2f[0, 0], 1)
            self.assertEqual(m2f[0, 1], 2)
            self.assertEqual(m2f[1, 0], 3)
            self.assertEqual(m2f[1, 1], 4)

        m2f = uLib.Math.Matrix2f()
        m2f[0, 0] = 1
        m2f[0, 1] = 2
        m2f[1, 0] = 3
        m2f[1, 1] = 4
        check_1234(m2f)

        m2f = uLib.Math.Matrix2f([1, 2, 3, 4])
        check_1234(m2f)

        # m2f = uLib.Math.Matrix2f([[1, 2], [3, 4]])
        # check_1234(m2f)

        m2f = uLib.Math.Matrix2f(np.array([[1, 2], [3, 4]]))
        check_1234(m2f)

    def test_vector2(self):
        v2f = uLib.Math.Vector2f()
        v2f[0] = 1
        v2f[1] = 2
        self.assertEqual(v2f[0], 1)
        self.assertEqual(v2f[1], 2)
        
        v2f = uLib.Math.Vector2f([1, 2])
        self.assertEqual(v2f[0], 1)
        self.assertEqual(v2f[1], 2)
        
        v2f = uLib.Math.Vector2f(np.array([1, 2]))
        self.assertEqual(v2f[0], 1)
        self.assertEqual(v2f[1], 2)
    
    def test_vector3(self):
        v3f = uLib.Math.Vector3f()
        v3f[0] = 1
        v3f[1] = 2
        v3f[2] = 3
        self.assertEqual(v3f[0], 1)
        self.assertEqual(v3f[1], 2)
        self.assertEqual(v3f[2], 3)
        
        v3f = uLib.Math.Vector3f([1, 2, 3])
        self.assertEqual(v3f[0], 1)
        self.assertEqual(v3f[1], 2)
        self.assertEqual(v3f[2], 3)
        
        v3f = uLib.Math.Vector3f(np.array([1, 2, 3]))
        self.assertEqual(v3f[0], 1)
        self.assertEqual(v3f[1], 2)
        self.assertEqual(v3f[2], 3)
    

class TestMathGeometry(unittest.TestCase):
    def test_geometry(self):
        Geo = uLib.Math.Geometry()
        
        Geo.SetPosition([1, 1, 1])
        
        pt = Geo.GetLocalPoint([2, 3, 2, 1])
        wp = Geo.GetWorldPoint(pt)
        
        self.assertTrue(vector4f0(wp, [2, 3, 2, 1]))
        
        Geo.Scale([2, 2, 2])
        wp = Geo.GetWorldPoint([1, 1, 1, 1])
        self.assertTrue(vector4f0(wp, [3, 3, 3, 1]))
        
class TestMathContainerBox(unittest.TestCase):
    def test_container_box_local(self):
        Cnt = uLib.Math.ContainerBox()
        Cnt.SetOrigin([-1, -1, -1])
        Cnt.SetSize([2, 2, 2])
        
        size = Cnt.GetSize()
        self.assertTrue(np.allclose(size, [2, 2, 2]))

    def test_container_box_global(self):
        Box = uLib.Math.ContainerBox()
        Box.SetPosition([1, 1, 1])
        Box.SetSize([2, 2, 2])
        
        pt = Box.GetLocalPoint([2, 3, 2, 1])
        wp = Box.GetWorldPoint(pt)
        self.assertTrue(vector4f0(wp, [2, 3, 2, 1]))

class TestMathStructuredGrid(unittest.TestCase):
    def test_structured_grid(self):
        grid = uLib.Math.StructuredGrid([10, 10, 10])
        grid.SetSpacing([1, 1, 1])
        
        spacing = grid.GetSpacing()
        self.assertTrue(np.allclose(spacing, [1, 1, 1]))

class TestMathAccumulator(unittest.TestCase):
    def test_accumulator_mean(self):
        acc = uLib.Math.Accumulator_Mean_f()
        acc(10.0)
        acc(20.0)
        self.assertAlmostEqual(acc(), 15.0)

class TestMathNewTypes(unittest.TestCase):
    def test_eigen_vectors(self):
        v1f = uLib.Math.Vector1f()
        v3d = uLib.Math.Vector3d()
        m4f = uLib.Math.Matrix4f()
        self.assertIsNotNone(v1f)
        self.assertIsNotNone(v3d)
        self.assertIsNotNone(m4f)

    def test_ulib_vectors(self):
        vi = uLib.Math.Vector_i()
        vi.append(1)
        vi.append(2)
        self.assertEqual(len(vi), 2)
        self.assertEqual(vi[0], 1)
        self.assertEqual(vi[1], 2)
        
        vf = uLib.Math.Vector_f()
        vf.append(1.5)
        self.assertAlmostEqual(vf[0], 1.5)

    def test_homogeneous(self):
        p = uLib.Math.HPoint3f(1.0, 2.0, 3.0)
        v = uLib.Math.HVector3f(0.0, 1.0, 0.0)
        self.assertIsNotNone(p)
        self.assertIsNotNone(v)

    def test_vox_image(self):
        img = uLib.Math.VoxImage([2, 2, 2])
        self.assertEqual(img.GetDims()[0], 2)
        img.SetValue([0, 0, 0], 10.5)
        # Note: GetValue returns float, and there might be internal scaling (1.E-6 observed in code)
        # Actually in VoxImage.h: GetValue(id) returns At(id).Value
        # SetValue(id, value) sets At(id).Value = value
        self.assertAlmostEqual(img.GetValue([0, 0, 0]), 10.5)

    def test_quad_mesh(self):
        mesh = uLib.Math.QuadMesh()
        mesh.AddPoint(uLib.Math.Vector3f([0, 0, 0]))
        mesh.AddPoint(uLib.Math.Vector3f([1, 0, 0]))
        mesh.AddPoint(uLib.Math.Vector3f([1, 1, 0]))
        mesh.AddPoint(uLib.Math.Vector3f([0, 1, 0]))
        
        mesh.AddQuad(uLib.Math.Vector4i([0, 1, 2, 3]))
        
        self.assertEqual(len(mesh.Points()), 4)
        self.assertEqual(len(mesh.Quads()), 1)

class TestMathVoxRaytracer(unittest.TestCase):
    def test_raytracer(self):
        grid = uLib.Math.StructuredGrid([10, 10, 10])
        grid.SetSpacing([1, 1, 1])
        grid.SetOrigin([0, 0, 0])
        
        rt = uLib.Math.VoxRaytracer(grid)
        self.assertIsNotNone(rt)
        
        # Test TraceBetweenPoints
        p1 = np.array([0.5, 0.5, -1.0, 1.0], dtype=np.float32)
        p2 = np.array([0.5, 0.5, 11.0, 1.0], dtype=np.float32)
        data = rt.TraceBetweenPoints(p1, p2)
        
        self.assertGreater(data.Count(), 0)
        self.assertAlmostEqual(data.TotalLength(), 10.0)
        
        # Check elements
        elements = data.Data()
        for i in range(data.Count()):
            self.assertGreaterEqual(elements[i].vox_id, 0)
            self.assertGreater(elements[i].L, 0)

    def test_ray_data(self):
        data = uLib.Math.VoxRaytracerRayData()
        data.SetCount(10)
        data.SetTotalLength(5.5)
        self.assertEqual(data.Count(), 10)
        self.assertAlmostEqual(data.TotalLength(), 5.5)

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