"""运行：python3 examples/check_rotation.py，全部检查使用 Python 标准库。"""

import math
import random
import unittest

from rotation import apply, exp_so3, identity, log_so3, mul, skew
from rotation import step_body, step_world, transpose


def determinant(r):
    return (r[0]*(r[4]*r[8]-r[5]*r[7])
            - r[1]*(r[3]*r[8]-r[5]*r[6])
            + r[2]*(r[3]*r[7]-r[4]*r[6]))


def unit(v):
    length = math.sqrt(sum(x*x for x in v))
    return [x/length for x in v]


class RotationChecks(unittest.TestCase):
    def close(self, actual, expected, tolerance=2e-11):
        self.assertEqual(len(actual), len(expected))
        self.assertLessEqual(max(abs(a-b) for a, b in zip(actual, expected)), tolerance)

    def test_known_rotations(self):
        self.close(exp_so3([0, 0, 0]), identity())
        rz = exp_so3([0, 0, math.pi/2])
        self.close(rz, [0, -1, 0, 1, 0, 0, 0, 0, 1])
        self.close(apply(rz, [2, -1, 3]), [1, 2, 3])
        cyclic = exp_so3([2*math.pi/(3*math.sqrt(3))]*3)
        self.close(cyclic, [0, 0, 1, 1, 0, 0, 0, 1, 0])
        rx = exp_so3([math.pi/2, 0, 0])
        ry = exp_so3([0, math.pi/2, 0])
        self.close(apply(mul(ry, rx), [0, 0, 1]), [0, -1, 0])
        self.close(apply(mul(rx, ry), [0, 0, 1]), [1, 0, 0])

    def test_invariants_and_roundtrip(self):
        rng = random.Random(20261005)
        angles = [0, 1e-12, 1e-8, 1e-4, 0.7, 2.5,
                  math.pi-1e-4, math.pi-1e-6, math.pi-1e-10,
                  math.pi-1e-14, math.pi, math.pi+1e-10, 4.8]
        angles += [rng.uniform(0, 2*math.pi) for _ in range(12)]
        axes = [[1, 0, 0], [0, -1, 0], [1, 1, 1], [-1, 1, 0]]
        axes += [[rng.uniform(-1, 1) for _ in range(3)] for _ in range(16)]
        for axis in axes:
            for angle in angles:
                with self.subTest(axis=axis, angle=angle):
                    phi = [angle*x for x in unit(axis)]
                    r = exp_so3(phi)
                    self.close(mul(transpose(r), r), identity())
                    self.assertAlmostEqual(determinant(r), 1.0, places=12)
                    recovered = log_so3(r)
                    self.close(exp_so3(recovered), r)
                    self.assertLessEqual(math.sqrt(sum(x*x for x in recovered)), math.pi+1e-14)
                    if angle < math.pi-1e-5:
                        self.close(recovered, phi)

    def test_half_turn_axis_convention(self):
        for axis in [[-1, 0, 0], [0, -1, 0], [0, 0, -1], [-2, 1, 1], [1, -3, 2]]:
            phi = log_so3(exp_so3([math.pi*x for x in unit(axis)]))
            k = max(range(3), key=lambda i: abs(phi[i]))
            self.assertGreater(phi[k], 0.0)
            self.assertAlmostEqual(math.sqrt(sum(x*x for x in phi)), math.pi, places=13)
        # 两个并列的最大分量，以第一个为正。
        self.close(log_so3([0, -1, 0, -1, 0, 0, 0, 0, -1]),
                   [math.pi/math.sqrt(2), -math.pi/math.sqrt(2), 0])

    def test_same_motion_in_two_frames(self):
        r = exp_so3([0.6, -0.8, 0.3])
        omega_body = [1.2, -0.4, 0.7]
        omega_world = apply(r, omega_body)
        for h in [0, 1e-5, 0.02, 0.8]:
            self.close(step_body(r, omega_body, h), step_world(r, omega_world, h))

    def test_finite_difference_derivative(self):
        r = exp_so3([0.4, -0.7, 0.9])
        omega = [0.3, 1.1, -0.8]
        epsilon = 1e-6
        for step, expected in [(step_body, mul(r, skew(omega))),
                               (step_world, mul(skew(omega), r))]:
            plus, minus = step(r, omega, epsilon), step(r, omega, -epsilon)
            derivative = [(a-b)/(2*epsilon) for a, b in zip(plus, minus)]
            self.close(derivative, expected, tolerance=2e-9)

    def test_euler_orthogonality_drift(self):
        r = exp_so3([0.3, 0.5, -0.2])
        omega, h = [0.7, -1.2, 0.8], 0.1
        k = skew(omega)
        euler = mul(r, [i+h*x for i, x in zip(identity(), k)])
        gram = mul(transpose(euler), euler)
        expected = [i+h*h*x for i, x in zip(identity(), mul(transpose(k), k))]
        self.close(gram, expected)
        self.assertGreater(max(abs(a-b) for a, b in zip(gram, identity())), 0.01)
        exact = step_body(r, omega, h)
        self.close(mul(transpose(exact), exact), identity())


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