#!/usr/bin/python3
import random
import numpy as np
import numpy.matlib as npmatlib
import numpy.linalg as nplinalg
from math import sin, cos, sqrt

def random_rotation(seed=None):
    if seed is not None:
        rng = random.Random(seed)
    else:
        rng = random.Random()

    # Generate an (not too short) axis and angle
    n = 0.0
    while n < 0.1:
        t = rng.random() * 3.14159265358979323846
        x = rng.random()
        y = rng.random()
        z = rng.random()
        n = x*x + y*y + z*z
    # and scale the axis to unit length.
    n = sqrt(n)
    x = x / n
    y = y / n
    z = z / n

    # Return rotation matrix for this axis and angle.
    s = sin(t)
    c = cos(t)
    u = 1.0 - c
    R = np.array([[ c + x*x*u,   x*y*u - z*s, x*z*u + y*s ],
                  [ y*x*u + z*s, c + y*y*u,   y*z*u - x*s ],
                  [ z*x*u - y*s, z*y*u + x*s, c + z*z*u   ]])
    return np.asmatrix(R)


if __name__ == '__main__':
    # We'll apply a "secret" random rotation S,
    S = random_rotation()
    print("Actual rotation matrix used:")
    print("    [ %9.6f %9.6f %9.6f ]" % (S[0,0], S[0,1], S[0,2]))
    print("    [ %9.6f %9.6f %9.6f ]" % (S[1,0], S[1,1], S[1,2]))
    print("    [ %9.6f %9.6f %9.6f ]" % (S[2,0], S[2,1], S[2,2]))
    print("")

    # to a hundred random 3D points b,
    b = npmatlib.rand(100, 3)

    # as a:
    a = b.dot(S)

    # To solve for S knowing only a and b, we first build B:
    B = np.asmatrix(npmatlib.zeros((3, 3)))
    for i in range(0, len(b)):
        # B += np.dot((b[i]).T, (a[i]))
        for row in range(0, 3):
            for col in range(0, 3):
                B[row, col] += b[i,row] * a[i,col]


    # Then we do Singular Value Decomposition on B.
    u, s, vh = nplinalg.svd(B, full_matrices=True)

    # Next, we construct T.
    T = npmatlib.eye(3)
    T[2,2] = nplinalg.det(u) * nplinalg.det(vh.T)

    # Finally, we construct the solution R.
    R = np.asmatrix(nplinalg.multi_dot([u, T, vh]))
    print("Rotation matrix found:")
    print("    [ %9.6f %9.6f %9.6f ]" % (R[0,0], R[0,1], R[0,2]))
    print("    [ %9.6f %9.6f %9.6f ]" % (R[1,0], R[1,1], R[1,2]))
    print("    [ %9.6f %9.6f %9.6f ]" % (R[2,0], R[2,1], R[2,2]))
    print("")

    # For the output, we also print the rotated coordinates.
    r = b.dot(S)

    print("Before   After   UsingSolution")
    for i in range(0, len(b)):
        print("%9.6f %9.6f %9.6f   %9.6f %9.6f %9.6f   %9.6f %9.6f %9.6f" %
              (b[i,0],b[i,1],b[i,2], a[i,0],a[i,1],a[i,2], r[i,0],r[i,1],r[i,2]))