import numpy as np
from math import cos, sin


def create_rx(unit_axis: np.ndarray) -> np.ndarray:
    (a, b, c) = unit_axis.tolist()
    d = (b ** 2 + c ** 2) ** 0.5
    return np.array(
        [[1, 0, 0, 0], [0, c / d, -b / d, 0], [0, b / d, c / d, 0], [0, 0, 0, 1]]
    )


def create_ry(unit_axis: np.ndarray) -> np.ndarray:
    (a, b, c) = unit_axis.tolist()
    d = (b ** 2 + c ** 2) ** 0.5
    return np.array([[d, 0, -a, 0], [0, 1, 0, 0], [a, 0, d, 0], [0, 0, 0, 1]])


def create_rz(angel: float) -> np.ndarray:
    return np.array(
        [
            [cos(angel), -sin(angel), 0, 0],
            [sin(angel), cos(angel), 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ]
    )


def create_t(axis: np.ndarray) -> np.ndarray:
    (x1, y1, z1) = axis.tolist()
    return np.array([[1, 0, 0, -x1], [0, 1, 0, -y1], [0, 0, 1, -z1], [0, 0, 0, 1]])


def rotation(angel: float, axis: np.ndarray, points: np.ndarray):
    # 依据《计算机图形学》第4版 9.28公式
    unit_axis = axis / np.linalg.norm(axis)
    rz = create_rz(angel)
    rx = create_rx(unit_axis)
    ry = create_ry(unit_axis)
    t = create_t(axis)
    t_i = np.linalg.inv(t)
    rx_i = np.linalg.inv(rx)
    ry_i = np.linalg.inv(ry)
    r_transform = t_i.dot(rx_i.dot(ry_i.dot(rz.dot(ry.dot(rx.dot(t))))))
    expand_point = np.concatenate((points, np.ones((points.shape[0], 1))), axis=1)
    return r_transform.dot(expand_point.T)