"""Analytic two-link dynamics and mass-weighted acceleration projection.

The shared rods move in world x-z, with positive axes along -y.
q2 is relative to the first rod; scalar joint velocities equal q derivatives.
Gravity is collected in bias, while applied generalized force is separate.
Matrices below are flat row-major arrays.
"""
import math
from functools import lru_cache
from kinematics import model_data
from inertia import box_inertia, dot


@lru_cache(maxsize=1)
def parameters():
    data = model_data()
    l1, l2 = data["lengths"]
    m1, m2 = data["masses"]
    r1, r2 = l1/2, l2/2
    i1, i2 = [box_inertia(m, [*data["cross_section"],l])[4]
              for m,l in zip(data["masses"],data["lengths"])]
    return {"l1":l1,"l2":l2,"m1":m1,"m2":m2,"r1":r1,"r2":r2,
            "i1":i1,"i2":i2,"g":-data["gravity"][2],
            "A":i1+m1*r1*r1+i2+m2*(l1*l1+r2*r2),
            "B":m2*l1*r2,"D":i2+m2*r2*r2}


def two_link_mass(q):
    p = parameters()
    coupling = p["B"]*math.cos(q[1])
    return [p["A"]+2*coupling,p["D"]+coupling,
            p["D"]+coupling,p["D"]]


def two_link_bias(q, velocity, gravity=None):
    p = parameters()
    g = p["g"] if gravity is None else gravity
    h = p["B"]*math.sin(q[1])
    v1, v2 = velocity
    g2 = g*p["m2"]*p["r2"]*math.sin(q[0]+q[1])
    g1 = g*(p["m1"]*p["r1"]+p["m2"]*p["l1"])*math.sin(q[0])+g2
    return [-h*(2*v1*v2+v2*v2)+g1,h*v1*v1+g2]


def two_link_energy(q, velocity, gravity=None):
    p = parameters()
    g = p["g"] if gravity is None else gravity
    m = two_link_mass(q)
    v1, v2 = velocity
    kinetic = (m[0]*v1*v1+2*m[1]*v1*v2+m[3]*v2*v2)/2
    potential = -g*((p["m1"]*p["r1"]+p["m2"]*p["l1"])*math.cos(q[0])
                    +p["m2"]*p["r2"]*math.cos(q[0]+q[1]))
    return kinetic, potential


def solve_spd(matrix, rhs):
    """Cholesky factorization, followed by two triangular solves; no inverse."""
    n = len(rhs)
    if len(matrix) != n*n:
        raise ValueError("Square matrix and right-hand side sizes must agree")
    lower = [0.0]*(n*n)
    for i in range(n):
        for j in range(i+1):
            if abs(matrix[n*i+j]-matrix[n*j+i]) > 1e-12:
                raise ValueError("Matrix must be symmetric")
            value = matrix[n*i+j]-sum(lower[n*i+k]*lower[n*j+k] for k in range(j))
            if i == j:
                if value <= 0:
                    raise ValueError("Matrix must be positive definite")
                lower[n*i+j] = math.sqrt(value)
            else:
                lower[n*i+j] = value/lower[n*j+j]
    intermediate, result = [0.0]*n, [0.0]*n
    for i in range(n):
        intermediate[i] = (rhs[i]-sum(lower[n*i+j]*intermediate[j] for j in range(i)))/lower[n*i+i]
    for i in reversed(range(n)):
        result[i] = (intermediate[i]-sum(lower[n*j+i]*result[j] for j in range(i+1,n)))/lower[n*i+i]
    return result


def two_link_acceleration(q, velocity, force=(0.0,0.0), gravity=None):
    c = two_link_bias(q,velocity,gravity)
    return solve_spd(two_link_mass(q),[x-y for x,y in zip(force,c)])


def step_two_link(state, dt, gravity=None, force=(0.0,0.0)):
    """One RK4 step; state is [q1,q2,v1,v2]."""
    def derivative(x):
        return x[2:]+two_link_acceleration(x[:2],x[2:],force,gravity)
    k1 = derivative(state)
    k2 = derivative([x+dt*k/2 for x,k in zip(state,k1)])
    k3 = derivative([x+dt*k/2 for x,k in zip(state,k2)])
    k4 = derivative([x+dt*k for x,k in zip(state,k3)])
    return [x+dt*(a+2*b+2*c+d)/6 for x,a,b,c,d in zip(state,k1,k2,k3,k4)]


def project_acceleration(mass, unconstrained, jacobian, target, regularizer=None):
    """Bilateral Gauss projection. Independent rows for hard constraints.

    For a positive diagonal regularizer, solve (A+R) lambda=b-J*au.
    An empty constraint list leaves acceleration unchanged.
    """
    n, rows = len(unconstrained), len(target)
    if len(jacobian) != n*rows:
        raise ValueError("Jacobian dimensions must match target and acceleration")
    if regularizer is not None and (len(regularizer)!=rows or min(regularizer,default=0)<0):
        raise ValueError("Regularizer must have nonnegative diagonal entries")
    directions = [solve_spd(mass,jacobian[n*i:n*(i+1)]) for i in range(rows)]
    a = [dot(jacobian[n*i:n*(i+1)],directions[j])
         for i in range(rows) for j in range(rows)]
    system = a[:]
    for i in range(rows):
        system[rows*i+i] += 0 if regularizer is None else regularizer[i]
    residual = [target[i]-dot(jacobian[n*i:n*(i+1)],unconstrained) for i in range(rows)]
    force = solve_spd(system,residual)
    correction = [sum(directions[i][j]*force[i] for i in range(rows)) for j in range(n)]
    result = [x+y for x,y in zip(unconstrained,correction)]
    generalized_force = [sum(jacobian[n*i+j]*force[i] for i in range(rows)) for j in range(n)]
    cost = sum(correction[i]*mass[n*i+j]*correction[j]
               for i in range(n) for j in range(n))/2
    return {"acceleration":result,"force":force,"generalized_force":generalized_force,
            "A":a,"correction":correction,"cost":cost}


if __name__ == "__main__":
    q, v = model_data()["default_angles"], [1,-.5]
    print("M =",two_link_mass(q))
    print("c =",two_link_bias(q,v))
    print("a =",two_link_acceleration(q,v))
    print("sliders =",project_acceleration([2,0,0,1],[2,-1],[1,-1],[0]))