1# -*- coding: utf-8 -*-
2from .Qt import QtGui
3from . import functions as fn
4from .Vector import Vector
5import numpy as np
6
7
8class Transform3D(QtGui.QMatrix4x4):
9    """
10    Extension of QMatrix4x4 with some helpful methods added.
11    """
12    def __init__(self, *args):
13        if len(args) == 1:
14            if isinstance(args[0], (list, tuple, np.ndarray)):
15                args = [x for y in args[0] for x in y]
16                if len(args) != 16:
17                    raise TypeError("Single argument to Transform3D must have 16 elements.")
18            elif isinstance(args[0], QtGui.QMatrix4x4):
19                args = list(args[0].copyDataTo())
20
21        QtGui.QMatrix4x4.__init__(self, *args)
22
23    def matrix(self, nd=3):
24        if nd == 3:
25            return np.array(self.copyDataTo()).reshape(4,4)
26        elif nd == 2:
27            m = np.array(self.copyDataTo()).reshape(4,4)
28            m[2] = m[3]
29            m[:,2] = m[:,3]
30            return m[:3,:3]
31        else:
32            raise Exception("Argument 'nd' must be 2 or 3")
33
34    def map(self, obj):
35        """
36        Extends QMatrix4x4.map() to allow mapping (3, ...) arrays of coordinates
37        """
38        if isinstance(obj, np.ndarray) and obj.shape[0] in (2,3):
39            if obj.ndim >= 2:
40                return fn.transformCoordinates(self, obj)
41            elif obj.ndim == 1:
42                v = QtGui.QMatrix4x4.map(self, Vector(obj))
43                return np.array([v.x(), v.y(), v.z()])[:obj.shape[0]]
44        elif isinstance(obj, (list, tuple)):
45            v = QtGui.QMatrix4x4.map(self, Vector(obj))
46            return type(obj)([v.x(), v.y(), v.z()])[:len(obj)]
47        else:
48            return QtGui.QMatrix4x4.map(self, obj)
49
50    def inverted(self):
51        inv, b = QtGui.QMatrix4x4.inverted(self)
52        return Transform3D(inv), b
53