"""第 2 章的姿态小内核：右手系、列向量、按行存储的九元素矩阵。

角度使用弧度。R 把刚体坐标变成世界坐标；各函数返回新列表。
log_so3 的输入为正交且行列式等于一的旋转矩阵。
"""

import math


def identity():
    """三阶单位矩阵。"""
    return [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]


def mul(a, b):
    """矩阵乘法 a b，最右侧的变换先作用。"""
    return [sum(a[3*i+k] * b[3*k+j] for k in range(3))
            for i in range(3) for j in range(3)]


def apply(r, v):
    """矩阵作用于三维列向量。"""
    return [sum(r[3*i+j] * v[j] for j in range(3)) for i in range(3)]


def transpose(r):
    return [r[3*j+i] for i in range(3) for j in range(3)]


def skew(v):
    """叉乘矩阵：apply(skew(v), p) 等于 v 与 p 的叉积。"""
    x, y, z = v
    return [0.0, -z, y, z, 0.0, -x, -y, x, 0.0]


def exp_so3(phi):
    """旋转向量转矩阵：向量方向给出轴，长度给出转角。"""
    theta2 = sum(x*x for x in phi)
    if theta2 < 1e-8:
        # 在零角附近直接计算两个系数的极限级数。
        a = 1.0 - theta2/6.0 + theta2*theta2/120.0
        b = 0.5 - theta2/24.0 + theta2*theta2/720.0
    else:
        theta = math.sqrt(theta2)
        a = math.sin(theta) / theta
        b = 2.0 * math.sin(theta/2.0)**2 / theta2
    k = skew(phi)
    k2 = mul(k, k)
    return [i + a*x + b*y for i, x, y in zip(identity(), k, k2)]


def log_so3(r):
    """矩阵转主值旋转向量，长度在 [0, pi] 内。

    半周处选轴的最大绝对分量为正；相等时优先较小下标。
    这个符号约定让同一输入得到确定的结果。
    """
    v = [(r[7]-r[5])/2.0, (r[2]-r[6])/2.0, (r[3]-r[1])/2.0]
    sine = math.sqrt(sum(x*x for x in v))
    cosine = max(-1.0, min(1.0, (r[0]+r[4]+r[8]-1.0)/2.0))
    theta = math.atan2(sine, cosine)
    if theta < 1e-6:
        theta2 = theta*theta
        scale = 1.0 + theta2/6.0 + 7.0*theta2*theta2/360.0
        return [scale*x for x in v]
    if math.pi - theta < 1e-5:
        # 对称部分给出轴的外积；最大对角元至少为三分之一。
        s = [(r[3*i+j] + r[3*j+i] - (2*cosine if i == j else 0.0))
             / (2.0*(1.0-cosine)) for i in range(3) for j in range(3)]
        k = max(range(3), key=lambda i: s[3*i+i])
        largest = math.sqrt(max(0.0, s[3*k+k]))
        axis = [s[3*i+k]/largest for i in range(3)]
        length = math.sqrt(sum(x*x for x in axis))
        axis = [x/length for x in axis]
        # 半周附近用反对称部分定向；浮点半周处沿用最大分量正号。
        if sine > 1e-14 and sum(x*y for x, y in zip(axis, v)) < 0:
            axis = [-x for x in axis]
        return [theta*x for x in axis]
    return [theta*x/sine for x in v]


def step_body(r, omega_body, h):
    """一步内保持体坐标角速度恒定，从右侧乘旋转增量。"""
    return mul(r, exp_so3([h*x for x in omega_body]))


def step_world(r, omega_world, h):
    """一步内保持世界坐标角速度恒定，从左侧乘旋转增量。"""
    return mul(exp_so3([h*x for x in omega_world]), r)