"""两条直线移动自由度的五阶段教学内核，采用 SI 单位。"""
from dataclasses import dataclass, field
from math import exp


def matvec(matrix, vector):
    return [sum(a*b for a, b in zip(row, vector)) for row in matrix]


def transpose_force(transmission, force):
    """执行器标量力经传动矩阵转置，成为广义力。"""
    return [sum(row[j]*f for row, f in zip(transmission, force))
            for j in range(len(transmission[0]))]


def solve2(matrix, rhs):
    """解可逆二阶系统；模型中的质量矩阵为正定对角矩阵。"""
    a, b = matrix[0]
    c, d = matrix[1]
    determinant = a*d-b*c
    if abs(determinant) < 1e-14:
        raise ValueError("线性系统的行列式接近零")
    return [(d*rhs[0]-b*rhs[1])/determinant,
            (a*rhs[1]-c*rhs[0])/determinant]


def equality_correction(mass, jacobian, smooth, target=0.0):
    """单行理想等式 J a = target；返回约束力、修正加速度、有效逆质量。"""
    direction = solve2(mass, jacobian)
    inverse_mass = sum(j*y for j, y in zip(jacobian, direction))
    if inverse_mass <= 0:
        raise ValueError("约束行需要一个非零的独立运动方向")
    multiplier = (target-sum(j*a for j, a in zip(jacobian, smooth)))/inverse_mass
    return multiplier, [a+y*multiplier for a, y in zip(smooth, direction)], inverse_mass


def filter_exact(activation, control, tau, h):
    """一步内控制恒定的一阶激活模型，tau 与 h 单位为秒。"""
    if tau <= 0 or h < 0:
        raise ValueError("时间常数为正，步长非负")
    return control+(activation-control)*exp(-h/tau)


@dataclass(frozen=True)
class Model:
    masses: tuple = (2.0, 1.0)
    transmission: tuple = ((1.0, 0.0), (0.5, 1.0))
    passive: tuple = (1.0, -1.0)
    bias: tuple = (2.0, 3.0)
    applied: tuple = (0.0, 4.0)
    jacobian: tuple = (-1.0, 1.0)


@dataclass
class Data:
    q: list = field(default_factory=lambda: [0.0, 0.0])
    v: list = field(default_factory=lambda: [0.2, 0.2])
    control: list = field(default_factory=lambda: [4.0, -2.0])
    time: float = 0.0
    cache: dict = field(default_factory=dict)


def forward(model, data, skip="none", constrained=True, target=0.0):
    """计算当前加速度，保持配置、速度和时间；skip 可取 none/position/velocity。"""
    if skip not in ("none", "position", "velocity"):
        raise ValueError("未知的阶段复用选项")
    if any(m <= 0 for m in model.masses):
        raise ValueError("质量为正")
    cache = data.cache
    key_position = (model, tuple(data.q))
    key_velocity = (key_position, tuple(data.v))
    if skip != "none" and cache.get("position_key") != key_position:
        raise ValueError("先计算当前配置对应的位置阶段")
    if skip == "velocity" and cache.get("velocity_key") != key_velocity:
        raise ValueError("先计算当前配置与速度对应的速度阶段")
    stages = []
    if skip == "none":
        cache.update(position_key=key_position, position=list(data.q),
                     mass=[[model.masses[0], 0.0], [0.0, model.masses[1]]],
                     length=matvec(model.transmission, data.q),
                     residual=sum(j*q for j, q in zip(model.jacobian, data.q)))
        stages.append("position")
    if skip != "velocity":
        cache.update(velocity_key=key_velocity,
                     speed=matvec(model.transmission, data.v),
                     passive=list(model.passive), bias=list(model.bias))
        stages.append("velocity")
    force = list(data.control)
    actuator = transpose_force(model.transmission, force)
    cache.update(force=force, actuator=actuator)
    stages.append("actuation")
    smooth_force = [p-c+f+u for p, c, f, u in
                    zip(cache["passive"], cache["bias"], model.applied, actuator)]
    smooth = solve2(cache["mass"], smooth_force)
    cache.update(smooth_force=smooth_force, smooth=smooth)
    stages.append("acceleration")
    multiplier, acceleration, inverse_mass = (0.0, list(smooth), 0.0)
    if constrained:
        multiplier, acceleration, inverse_mass = equality_correction(
            cache["mass"], model.jacobian, smooth, target)
    cache.update(multiplier=multiplier, acceleration=acceleration,
                 constraint=[j*multiplier for j in model.jacobian],
                 inverse_mass=inverse_mass, evaluated_at=data.time)
    stages.append("constraint")
    return stages


def step(model, data, h, constrained=True, target=0.0):
    """半隐式 Euler：保留这一步计算的派生量，推进 q、v、time。"""
    if h <= 0:
        raise ValueError("步长为正")
    forward(model, data, constrained=constrained, target=target)
    data.v = [v+h*a for v, a in zip(data.v, data.cache["acceleration"])]
    data.q = [q+h*v for q, v in zip(data.q, data.v)]
    data.time += h


if __name__ == "__main__":
    model, data = Model(), Data()
    forward(model, data)
    print("平滑加速度：", data.cache["smooth"])
    print("约束力与最终加速度：", data.cache["multiplier"], data.cache["acceleration"])
    step(model, data, 0.1)
    print("新配置与缓存配置：", data.q, data.cache["position"])