#!/usr/bin/env python
#
__author__    = 'Mike McCann'

__doc__ = '''
Test computing 'euler vector' axis and rotation rate by division by subsequent quaternions.
See https://github.com/ezag/pyeuclid and https://github.com/ezag/pyeuclid/blob/master/euclid.txt.
'''

import sys
import math
import unittest
import numpy as np
from BEDS import BEDS
from bed2x3d import BEDS_X3D
from argparse import Namespace
from euclid import Quaternion, Vector3

import pdb

class TestDivision(unittest.TestCase):

    def setUp(self):
        pass

    def testBasicDiv(self):
        '''Test Basic division of Quaternions
        '''
        q1 = Quaternion.new_rotate_axis(0, Vector3(1, 0, 0))
        q2 = Quaternion.new_rotate_axis(math.pi / 2, Vector3(1, 0, 0))
        q3 = Quaternion.new_rotate_axis(math.pi / 2, Vector3(0, 1, 0))
        q4 = Quaternion.new_rotate_axis(math.pi / 2, Vector3(0, 0, 1))

        dq = q2 * Quaternion.conjugated(q1)
        aa = dq.get_angle_axis()

        self.assertEqual(aa[0], math.pi / 2)

        self.assertEqual((q3 * Quaternion.conjugated(q1)).get_angle_axis()[0], math.pi / 2)
        self.assertEqual((q4 * Quaternion.conjugated(q1)).get_angle_axis()[0], math.pi / 2)

class TestBEDSprocess(unittest.TestCase):

    def setUp(self):
        '''Create a known set of rotations simulating data recorded on the BED
        '''
        self.quatList = []

        self.rx = []
        self.ry = []
        self.rz = []

        self.px = []
        self.py = []
        self.pz = []
        self.prot = []

        self.mx = [0]
        self.my = [0]
        self.mz = [0]
        self.drot = [0]

        # Angles 0.0 and 2 * np.pi are degenerate - start with a 90 deg rotation and rotate 90 deg more
        astart = np.pi / 2
        aend = np.pi

        # A sequence of makeRotations must start where the previous one left. For some reason X & Y are swapped beteen Quaternion & X3D
        # It'd be nice to connect these rotations so that they execute one after another, but my mind was already spinning around
        # enough to figure that out such that dr was consistent. For now the other rotations can be run individually by uncommenting them.
        self.makeRotation('AboutX', Quaternion.new_rotate_axis(astart, Vector3(0, 1, 0)), Quaternion.new_rotate_axis(aend, Vector3(0, 1, 0)))
        ##self.makeRotation('AboutY', Quaternion.new_rotate_axis(astart, Vector3(1, 0, 0)), Quaternion.new_rotate_axis(aend, Vector3(1, 0, 0)))
        ##self.makeRotation('AboutZ', Quaternion.new_rotate_axis(astart, Vector3(0, 0, 1)), Quaternion.new_rotate_axis(aend, Vector3(0, 0, 1)))

    def makeRotation(self, name, q1, q2, number=6):
        '''Generate test data of a rotation between q1 and q1
        '''
        print name + ': interpolating between ' + str(q1) + ' and ' + str(q2) + ' to make test data'
        drot = 0
        drot_sum = 0
        for i in range(number):
            q = Quaternion.new_interpolate(q1, q2, i / float(number-1))
            self.quatList.append((q.w, q.x, q.y, q.z))

            rx, ry, rz = q.get_euler()
            self.rx.append(rx)
            self.ry.append(ry)
            self.rz.append(rz)

            prot, quat = q.get_angle_axis()
            self.px.append(quat.x)
            self.py.append(quat.y)
            self.pz.append(quat.z)
            self.prot.append(prot)

            if i > 0:
                diff_rot = q * Quaternion.conjugated(last_q)
                drot, dquat = diff_rot.get_angle_axis()
                self.mx.append(dquat.x)
                self.my.append(dquat.y)
                self.mz.append(dquat.z)
                self.drot.append(drot)

            last_q = q

            drot_sum += drot
            fmtStr = '%2d. rx, ry, rz, drot, drot_sum = %6.2f %6.2f %6.2f %6.2f %6.2f'
            fmtStr += '   px, py, pz, prot = %6.2f %6.2f %6.2f %6.2f'
            print fmtStr % (i, rx, ry, rz, drot, drot_sum, quat.x, quat.y, quat.z, prot)

    def makeX3Dfile(self):
        '''Use bed2x3d module to write a .x3d for visual verification
        '''
        args = Namespace()
        args.verbose = 0
        args.input = []
        args.translate = False
        args.speedup = 1
        args.output = 'testBEDSprocess.x3d'

        x3d = BEDS_X3D(args, sys.argv)
        x3d.sensorType = 'Invensense'
        x3d.rateHz = 1.0
        x3d.s = np.arange(0.0, len(self.rx)/x3d.rateHz)
        x3d.quatList = self.quatList

        x3d.ax = np.zeros(len(self.rx))     # Dummy-up 0 accelerations for createX3D()
        x3d.ay = np.zeros(len(self.rx))
        x3d.az = np.zeros(len(self.rx))

        x3dText = x3d.createX3D()           # Calls BEDS processRotations()

        f = open(args.output, 'w')
        f.write(x3dText)
        f.close()

        print "Wrote file %s - open it in InstantReality Player, BS Contact, Xj3D, X3DOM, or other X3D browser." % args.output

    def testProcessRotation(self):
        '''Simulation of data from BED with a known rotation
        '''
        beds = BEDS()
        beds.args = Namespace()
        beds.args.verbose = 1

        beds.rateHz = 10.0
        beds.quatList = self.quatList

        print 'Testing BEDS processRotations...'
        beds.processRotations()

        print 'Euler angles'
        for rx1, rx2 in zip(beds.rxList, self.rx):
            self.assertEqual(rx1, rx2)

        for ry1, ry2 in zip(beds.ryList, self.ry):
            self.assertEqual(ry1, ry2)

        for rz1, rz2 in zip(beds.rzList, self.rz):
            self.assertEqual(rz1, rz2)

        print 'Euler vectors (angle_axis & X3D orientation)'
        for px1, px2 in zip(beds.pxList, self.px):
            self.assertEqual(px1, px2)

        for py1, py2 in zip(beds.pyList, self.py):
            self.assertEqual(py1, py2)

        for pz1, pz2 in zip(beds.pzList, self.pz):
            self.assertEqual(pz1, pz2)

        for prot1, prot2 in zip(beds.protList, self.prot):
            self.assertEqual(prot1, prot2)

        print 'Rotation difference'
        for dr1, dr2 in zip(beds.diffrotList, self.drot):
            self.assertEqual(dr1, dr2)

        # Make X3D file for visual verification
        self.makeX3Dfile()


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