text:
## Contains functions for simulating the delta robot or a single delta robot arm
using LinearAlgebra, StaticArrays

"""
generates initial conditions for delta arm IVP
"""
function delta_arm_ics(u0::Array{T, 1}, p::DeltaRobotParams) where {T<:Real}
    return u0
end

"""
generates initial conditions for delta robot IVP
"""
function delta_robot_ics(u0::Array{T, 1}, du0::Array{T, 1}, ddu0::Array{T, 1}) where {T<:Real}
    u0 = vcat(u0, du0, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    du0 = vcat(du0, ddu0, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
    diff_vars = [true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, false, false, false, false, false, false]
    return u0, du0, diff_vars
end

"""
delta_arm!(dx, d, p, t) defines the differential equation for single
delta robot arm with damping
"""
function delta_arm!(dx, x, p, t)
    # states of plant
    phi1 = x[1];
    phi2 = x[2];
    phi3 = x[3];

    phi1dot = x[4];
    phi2dot = x[5];
    phi3dot = x[6];

    # model params
    m1 = p.m1;
    m2 = p.m2;
    m3 = p.m3;
    l1 = p.l1;
    l2 = p.l2;
    g = p.g;

    # inertia matrix
    M1 = [(l1^2*(m1 + 3*m2 + m3))/3; (l1*l2*(sin(phi1)*sin(phi2) + cos(phi1)*cos(phi2)*cos(phi3))*(3*m2 + 2*m3))/6; -(l1*l2*cos(phi1)*sin(phi2)*sin(phi3)*(3*m2 + 2*m3))/6]';
    M2 = [(l1*l2*(sin(phi1)*sin(phi2) + cos(phi1)*cos(phi2)*cos(phi3))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M3 = [-(l1*l2*cos(phi1)*sin(phi2)*sin(phi3)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(phi2)^2 - 1))/3]';
    M = [M1; M2; M3];

    # coriolis & gravitational matrices (basically external forces)
    b = [
        (g*l1*m1*cos(phi1))/2 + g*l1*m2*cos(phi1) + (g*l1*m3*cos(phi1))/3 - (l1*l2*m2*phi2dot^2*cos(phi2)*sin(phi1))/2 - (l1*l2*m3*phi2dot^2*cos(phi2)*sin(phi1))/3 + (l1*l2*m2*phi2dot^2*cos(phi1)*cos(phi3)*sin(phi2))/2 + (l1*l2*m2*phi3dot^2*cos(phi1)*cos(phi3)*sin(phi2))/2 + (l1*l2*m3*phi2dot^2*cos(phi1)*cos(phi3)*sin(phi2))/3 + (l1*l2*m3*phi3dot^2*cos(phi1)*cos(phi3)*sin(phi2))/3 + l1*l2*m2*phi2dot*phi3dot*cos(phi1)*cos(phi2)*sin(phi3) + (2*l1*l2*m3*phi2dot*phi3dot*cos(phi1)*cos(phi2)*sin(phi3))/3;
        (l2*(l2*m2*phi3dot^2*sin(2*phi2) + l2*m3*phi3dot^2*sin(2*phi2) + 3*g*m2*cos(phi2)*cos(phi3) + 2*g*m3*cos(phi2)*cos(phi3) - 3*l1*m2*phi1dot^2*cos(phi1)*sin(phi2) - 2*l1*m3*phi1dot^2*cos(phi1)*sin(phi2) + 3*l1*m2*phi1dot^2*cos(phi2)*cos(phi3)*sin(phi1) + 2*l1*m3*phi1dot^2*cos(phi2)*cos(phi3)*sin(phi1)))/6;
        -(l2*(3*g*m2*sin(phi2)*sin(phi3) + 2*g*m3*sin(phi2)*sin(phi3) + 2*l2*m2*phi2dot*phi3dot*sin(2*phi2) + 2*l2*m3*phi2dot*phi3dot*sin(2*phi2) + 3*l1*m2*phi1dot^2*sin(phi1)*sin(phi2)*sin(phi3) + 2*l1*m3*phi1dot^2*sin(phi1)*sin(phi2)*sin(phi3)))/6
    ];

    # Damping matrix
    D = zeros(3);
    D[1] = p.d1;
    D[2] = p.d2;
    D[3] = p.d3;

    # differential equation
    dx[1:3] = x[4:6];
    dx[4:6] = inv(M) * (b .- D .* x[4:6]);
end

"""
L_delta_robot!(out, dx, x, p, t) is the DAE update equation for the entire
delta robot under the lagrangian formulation
"""
function L_delta_robot!(out, dx, x, p, t)
    # states of plant
    q11 = x[1];
    q12 = x[2];
    q13 = x[3];
    q21 = x[4];
    q22 = x[5];
    q23 = x[6];
    q31 = x[7];
    q32 = x[8];
    q33 = x[9];
    q11dot = x[10];
    q12dot = x[11];
    q13dot = x[12];
    q21dot = x[13];
    q22dot = x[14];
    q23dot = x[15];
    q31dot = x[16];
    q32dot = x[17];
    q33dot = x[18];
    lambda1 = x[19];
    lambda2 = x[20];
    lambda3 = x[21];
    lambda4 = x[22];
    lambda5 = x[23];
    lambda6 = x[24];

    # model params
    m1 = p.m1;
    m2 = p.m2;
    m3 = p.m3;
    l1 = p.l1;
    l2 = p.l2;
    g = p.g;

    # 'D' elements
    M1 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6]';
    M2 = [(l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M3 = [-(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q12)^2 - 1))/3]';
    Mtop = [M1; M2; M3];
    M4 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6]';
    M5 = [(l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M6 = [-(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q22)^2 - 1))/3]';
    Mmid = [M4; M5; M6];
    M7 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6]';
    M8 = [(l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M9 = [-(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q32)^2 - 1))/3]';
    Mbot = [M7; M8; M9];
    invMtop = inv(Mtop);
    invMmid = inv(Mmid);
    invMbot = inv(Mbot);
    invM1 = [invMtop[1,:]; 0; 0; 0; 0; 0; 0]';
    invM2 = [invMtop[2,:]; 0; 0; 0; 0; 0; 0]';
    invM3 = [invMtop[3,:]; 0; 0; 0; 0; 0; 0]';
    invM4 = [0; 0; 0; invMmid[1,:]; 0; 0; 0]';
    invM5 = [0; 0; 0; invMmid[2,:]; 0; 0; 0]';
    invM6 = [0; 0; 0; invMmid[3,:]; 0; 0; 0]';
    invM7 = [0; 0; 0; 0; 0; 0; invMbot[1,:]]';
    invM8 = [0; 0; 0; 0; 0; 0; invMbot[2,:]]';
    invM9 = [0; 0; 0; 0; 0; 0; invMbot[3,:]]';
    invM = [invM1; invM2; invM3; invM4; invM5; invM6; invM7; invM8; invM9]

    # External Forces
    b = [
        (l1*(6*lambda3*cos(q11) + 6*lambda6*cos(q11) - 6*lambda2*sin(q11) - 6*lambda5*sin(q11) + 3*g*m1*cos(q11) + 6*g*m2*cos(q11) + 2*g*m3*cos(q11) - 3*l2*m2*q12dot^2*cos(q12)*sin(q11) - 2*l2*m3*q12dot^2*cos(q12)*sin(q11) + 3*l2*m2*q12dot^2*cos(q11)*cos(q13)*sin(q12) + 3*l2*m2*q13dot^2*cos(q11)*cos(q13)*sin(q12) + 2*l2*m3*q12dot^2*cos(q11)*cos(q13)*sin(q12) + 2*l2*m3*q13dot^2*cos(q11)*cos(q13)*sin(q12) + 6*l2*m2*q12dot*q13dot*cos(q11)*cos(q12)*sin(q13) + 4*l2*m3*q12dot*q13dot*cos(q11)*cos(q12)*sin(q13)))/6;
        (l2*(6*lambda3*cos(q12)*cos(q13) - 6*lambda5*sin(q12) - 6*lambda2*sin(q12) + 6*lambda6*cos(q12)*cos(q13) + 6*lambda1*cos(q12)*sin(q13) + 6*lambda4*cos(q12)*sin(q13) + l2*m2*q13dot^2*sin(2*q12) + l2*m3*q13dot^2*sin(2*q12) + 3*g*m2*cos(q12)*cos(q13) + 2*g*m3*cos(q12)*cos(q13) - 3*l1*m2*q11dot^2*cos(q11)*sin(q12) - 2*l1*m3*q11dot^2*cos(q11)*sin(q12) + 3*l1*m2*q11dot^2*cos(q12)*cos(q13)*sin(q11) + 2*l1*m3*q11dot^2*cos(q12)*cos(q13)*sin(q11)))/6;
        -(l2*(6*lambda3*sin(q12)*sin(q13) - 6*lambda4*cos(q13)*sin(q12) - 6*lambda1*cos(q13)*sin(q12) + 6*lambda6*sin(q12)*sin(q13) + 3*g*m2*sin(q12)*sin(q13) + 2*g*m3*sin(q12)*sin(q13) + 2*l2*m2*q12dot*q13dot*sin(2*q12) + 2*l2*m3*q12dot*q13dot*sin(2*q12) + 3*l1*m2*q11dot^2*sin(q11)*sin(q12)*sin(q13) + 2*l1*m3*q11dot^2*sin(q11)*sin(q12)*sin(q13)))/6;
        (l1*(3*g*m1*cos(q21) - 3*lambda2*sin(q21) - 6*lambda3*cos(q21) + 6*g*m2*cos(q21) + 2*g*m3*cos(q21) + 3*3^(1/2)*lambda1*sin(q21) - 3*l2*m2*q22dot^2*cos(q22)*sin(q21) - 2*l2*m3*q22dot^2*cos(q22)*sin(q21) + 3*l2*m2*q22dot^2*cos(q21)*cos(q23)*sin(q22) + 3*l2*m2*q23dot^2*cos(q21)*cos(q23)*sin(q22) + 2*l2*m3*q22dot^2*cos(q21)*cos(q23)*sin(q22) + 2*l2*m3*q23dot^2*cos(q21)*cos(q23)*sin(q22) + 6*l2*m2*q22dot*q23dot*cos(q21)*cos(q22)*sin(q23) + 4*l2*m3*q22dot*q23dot*cos(q21)*cos(q22)*sin(q23)))/6;
        (l2*(3*3^(1/2)*lambda1*sin(q22) - 3*lambda2*sin(q22) - 6*lambda3*cos(q22)*cos(q23) + 3*lambda1*cos(q22)*sin(q23) + l2*m2*q23dot^2*sin(2*q22) + l2*m3*q23dot^2*sin(2*q22) + 3*g*m2*cos(q22)*cos(q23) + 2*g*m3*cos(q22)*cos(q23) + 3*3^(1/2)*lambda2*cos(q22)*sin(q23) - 3*l1*m2*q21dot^2*cos(q21)*sin(q22) - 2*l1*m3*q21dot^2*cos(q21)*sin(q22) + 3*l1*m2*q21dot^2*cos(q22)*cos(q23)*sin(q21) + 2*l1*m3*q21dot^2*cos(q22)*cos(q23)*sin(q21)))/6;
        -(l2*(3*g*m2*sin(q22)*sin(q23) - 6*lambda3*sin(q22)*sin(q23) - 3*3^(1/2)*lambda2*cos(q23)*sin(q22) - 3*lambda1*cos(q23)*sin(q22) + 2*g*m3*sin(q22)*sin(q23) + 2*l2*m2*q22dot*q23dot*sin(2*q22) + 2*l2*m3*q22dot*q23dot*sin(2*q22) + 3*l1*m2*q21dot^2*sin(q21)*sin(q22)*sin(q23) + 2*l1*m3*q21dot^2*sin(q21)*sin(q22)*sin(q23)))/6;
        (l1*(3*g*m1*cos(q31) - 3*lambda5*sin(q31) - 6*lambda6*cos(q31) + 6*g*m2*cos(q31) + 2*g*m3*cos(q31) - 3*3^(1/2)*lambda4*sin(q31) - 3*l2*m2*q32dot^2*cos(q32)*sin(q31) - 2*l2*m3*q32dot^2*cos(q32)*sin(q31) + 3*l2*m2*q32dot^2*cos(q31)*cos(q33)*sin(q32) + 3*l2*m2*q33dot^2*cos(q31)*cos(q33)*sin(q32) + 2*l2*m3*q32dot^2*cos(q31)*cos(q33)*sin(q32) + 2*l2*m3*q33dot^2*cos(q31)*cos(q33)*sin(q32) + 6*l2*m2*q32dot*q33dot*cos(q31)*cos(q32)*sin(q33) + 4*l2*m3*q32dot*q33dot*cos(q31)*cos(q32)*sin(q33)))/6;
        (l2*(3*lambda4*cos(q32)*sin(q33) - 3*3^(1/2)*lambda4*sin(q32) - 6*lambda6*cos(q32)*cos(q33) - 3*lambda5*sin(q32) + l2*m2*q33dot^2*sin(2*q32) + l2*m3*q33dot^2*sin(2*q32) + 3*g*m2*cos(q32)*cos(q33) + 2*g*m3*cos(q32)*cos(q33) - 3*3^(1/2)*lambda5*cos(q32)*sin(q33) - 3*l1*m2*q31dot^2*cos(q31)*sin(q32) - 2*l1*m3*q31dot^2*cos(q31)*sin(q32) + 3*l1*m2*q31dot^2*cos(q32)*cos(q33)*sin(q31) + 2*l1*m3*q31dot^2*cos(q32)*cos(q33)*sin(q31)))/6;
        -(l2*(3*3^(1/2)*lambda5*cos(q33)*sin(q32) - 6*lambda6*sin(q32)*sin(q33) - 3*lambda4*cos(q33)*sin(q32) + 3*g*m2*sin(q32)*sin(q33) + 2*g*m3*sin(q32)*sin(q33) + 2*l2*m2*q32dot*q33dot*sin(2*q32) + 2*l2*m3*q32dot*q33dot*sin(2*q32) + 3*l1*m2*q31dot^2*sin(q31)*sin(q32)*sin(q33) + 2*l1*m3*q31dot^2*sin(q31)*sin(q32)*sin(q33)))/6
    ];

    # Holonomic Constraint
    h = [
        l2*sin(q12)*sin(q13) - (3^(1/2)*(l1*cos(q21) + l2*cos(q22)))/2 + (l2*sin(q22)*sin(q23))/2;
        l1*cos(q11) + l2*cos(q12) + (l1*cos(q21))/2 + (l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2;
        l1*sin(q11) - l1*sin(q21) + l2*cos(q13)*sin(q12) - l2*cos(q23)*sin(q22);
        (3^(1/2)*(l1*cos(q31) + l2*cos(q32)))/2 + l2*sin(q12)*sin(q13) + (l2*sin(q32)*sin(q33))/2;
        l1*cos(q11) + l2*cos(q12) + (l1*cos(q31))/2 + (l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2;
        l1*sin(q11) - l1*sin(q31) + l2*cos(q13)*sin(q12) - l2*cos(q33)*sin(q32)
    ];

    # Jacobian of Holonomic Constraint w.r.t q
    H1 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), (3^(1/2)*l1*sin(q21))/2, (3^(1/2)*l2*sin(q22))/2 + (l2*cos(q22)*sin(q23))/2, (l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H2 = [-l1*sin(q11), -l2*sin(q12), 0, -(l1*sin(q21))/2, (3^(1/2)*l2*cos(q22)*sin(q23))/2 - (l2*sin(q22))/2, (3^(1/2)*l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H3 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), -l1*cos(q21), -l2*cos(q22)*cos(q23), l2*sin(q22)*sin(q23), 0, 0, 0]';
    H4 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), 0, 0, 0, -(3^(1/2)*l1*sin(q31))/2, (l2*cos(q32)*sin(q33))/2 - (3^(1/2)*l2*sin(q32))/2, (l2*cos(q33)*sin(q32))/2]';
    H5 = [-l1*sin(q11), -l2*sin(q12), 0, 0, 0, 0, -(l1*sin(q31))/2, - (l2*sin(q32))/2 - (3^(1/2)*l2*cos(q32)*sin(q33))/2, -(3^(1/2)*l2*cos(q33)*sin(q32))/2]';
    H6 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), 0, 0, 0, -l1*cos(q31), -l2*cos(q32)*cos(q33), l2*sin(q32)*sin(q33)]';
    H = [
        H1; H2; H3; H4; H5; H6
    ];

    # Time Derivative of Jacobian w.r.t of Holonomic Constraint
    Hdot1 = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), (3^(1/2)*l1*q21dot*cos(q21))/2, (l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*sin(q22)*sin(q23))/2 - (3^(1/2)*l2*cos(q22))/2), (l2*q22dot*cos(q22)*cos(q23))/2 - (l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    Hdot2 = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, -(l1*q21dot*cos(q21))/2, (3^(1/2)*l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2), (3^(1/2)*l2*q22dot*cos(q22)*cos(q23))/2 - (3^(1/2)*l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    Hdot3 = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), l1*q21dot*sin(q21), l2*q22dot*cos(q23)*sin(q22) + l2*q23dot*cos(q22)*sin(q23), l2*q22dot*cos(q22)*sin(q23) + l2*q23dot*cos(q23)*sin(q22), 0, 0, 0]';
    Hdot4 = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), 0, 0, 0, -(3^(1/2)*l1*q31dot*cos(q31))/2, (l2*q33dot*cos(q32)*cos(q33))/2 - q32dot*((l2*sin(q32)*sin(q33))/2 + (3^(1/2)*l2*cos(q32))/2), (l2*q32dot*cos(q32)*cos(q33))/2 - (l2*q33dot*sin(q32)*sin(q33))/2]';
    Hdot5 = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, 0, 0, 0, -(l1*q31dot*cos(q31))/2, - q32dot*((l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2) - (3^(1/2)*l2*q33dot*cos(q32)*cos(q33))/2, (3^(1/2)*l2*q33dot*sin(q32)*sin(q33))/2 - (3^(1/2)*l2*q32dot*cos(q32)*cos(q33))/2]';
    Hdot6 = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), 0, 0, 0, l1*q31dot*sin(q31), l2*q32dot*cos(q33)*sin(q32) + l2*q33dot*cos(q32)*sin(q33), l2*q32dot*cos(q32)*sin(q33) + l2*q33dot*cos(q33)*sin(q32)]';
    Hdot = [
        Hdot1; Hdot2; Hdot3; Hdot4; Hdot5; Hdot6
    ];

    D = zeros(9);
    D[1] = p.d1;
    D[2] = p.d2;
    D[3] = p.d3;
    D[4] = p.d1;
    D[5] = p.d2;
    D[6] = p.d3;
    D[7] = p.d1;
    D[8] = p.d2;
    D[9] = p.d3;

    invD_times_b = invD * (b .- D * x[10:18]); # invD_times_b = invD * b;
    z0 = h;
    z1 = H * x[10:18];
    z2 = Hdot * x[10:18] .+ H * invD_times_b;
    z = z0 + z1 + z2;

    out[1:9] = x[10:18] .- dx[1:9]; # first derivative residual
    out[10:18] = invD_times_b .- dx[10:18]; # second derivative residual
    out[19:24] = z; # index reduced holonomic constraint
end

"""
L_delta_robot_base!(out, dx, x, p, t) DAE update equation for entire delta robot
under Lagrange formulation accounting for both base and platform radii
"""
function L_delta_robot_base!(out, dx, x, p, t)
    # states of plant
    q11 = x[1];
    q12 = x[2];
    q13 = x[3];
    q21 = x[4];
    q22 = x[5];
    q23 = x[6];
    q31 = x[7];
    q32 = x[8];
    q33 = x[9];
    q11dot = x[10];
    q12dot = x[11];
    q13dot = x[12];
    q21dot = x[13];
    q22dot = x[14];
    q23dot = x[15];
    q31dot = x[16];
    q32dot = x[17];
    q33dot = x[18];
    lambda1 = x[19];
    lambda2 = x[20];
    lambda3 = x[21];
    lambda4 = x[22];
    lambda5 = x[23];
    lambda6 = x[24];

    # model params
    m1 = p.m1;
    m2 = p.m2;
    m3 = p.m3;
    l1 = p.l1;
    l2 = p.l2;
    g = p.g;
    r_base = p.rb;
    r_platform = p.rp

    # inertia matrix
    M1 = [(l1^ 2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6]';
    M2 = [(l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M3 = [-(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q12)^2 - 1))/3]';
    Mtop = [M1; M2; M3];

    M4 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6]';
    M5 = [(l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M6 = [-(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q22)^2 - 1))/3]';
    Mmid = [M4; M5; M6];

    M7 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6]';
    M8 = [(l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M9 = [-(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q32)^2 - 1))/3]';
    Mbot = [M7; M8; M9];

    invMtop = inv(Mtop);
    invMmid = inv(Mmid);
    invMbot = inv(Mbot);
    invM1 = [invMtop[1,:]; 0; 0; 0; 0; 0; 0]';
    invM2 = [invMtop[2,:]; 0; 0; 0; 0; 0; 0]';
    invM3 = [invMtop[3,:]; 0; 0; 0; 0; 0; 0]';
    invM4 = [0; 0; 0; invMmid[1,:]; 0; 0; 0]';
    invM5 = [0; 0; 0; invMmid[2,:]; 0; 0; 0]';
    invM6 = [0; 0; 0; invMmid[3,:]; 0; 0; 0]';
    invM7 = [0; 0; 0; 0; 0; 0; invMbot[1,:]]';
    invM8 = [0; 0; 0; 0; 0; 0; invMbot[2,:]]';
    invM9 = [0; 0; 0; 0; 0; 0; invMbot[3,:]]';
    invM = [invM1; invM2; invM3; invM4; invM5; invM6; invM7; invM8; invM9]

    b = [
        (l1*(6*lambda3*cos(q11) + 6*lambda6*cos(q11) - 6*lambda2*sin(q11) - 6*lambda5*sin(q11) + 3*g*m1*cos(q11) + 6*g*m2*cos(q11) + 2*g*m3*cos(q11) - 3*l2*m2*q12dot^2*cos(q12)*sin(q11) - 2*l2*m3*q12dot^2*cos(q12)*sin(q11) + 3*l2*m2*q12dot^2*cos(q11)*cos(q13)*sin(q12) + 3*l2*m2*q13dot^2*cos(q11)*cos(q13)*sin(q12) + 2*l2*m3*q12dot^2*cos(q11)*cos(q13)*sin(q12) + 2*l2*m3*q13dot^2*cos(q11)*cos(q13)*sin(q12) + 6*l2*m2*q12dot*q13dot*cos(q11)*cos(q12)*sin(q13) + 4*l2*m3*q12dot*q13dot*cos(q11)*cos(q12)*sin(q13)))/6;
        (l2*(6*lambda3*cos(q12)*cos(q13) - 6*lambda5*sin(q12) - 6*lambda2*sin(q12) + 6*lambda6*cos(q12)*cos(q13) + 6*lambda1*cos(q12)*sin(q13) + 6*lambda4*cos(q12)*sin(q13) + l2*m2*q13dot^2*sin(2*q12) + l2*m3*q13dot^2*sin(2*q12) + 3*g*m2*cos(q12)*cos(q13) + 2*g*m3*cos(q12)*cos(q13) - 3*l1*m2*q11dot^2*cos(q11)*sin(q12) - 2*l1*m3*q11dot^2*cos(q11)*sin(q12) + 3*l1*m2*q11dot^2*cos(q12)*cos(q13)*sin(q11) + 2*l1*m3*q11dot^2*cos(q12)*cos(q13)*sin(q11)))/6;
        -(l2*(6*lambda3*sin(q12)*sin(q13) - 6*lambda4*cos(q13)*sin(q12) - 6*lambda1*cos(q13)*sin(q12) + 6*lambda6*sin(q12)*sin(q13) + 3*g*m2*sin(q12)*sin(q13) + 2*g*m3*sin(q12)*sin(q13) + 2*l2*m2*q12dot*q13dot*sin(2*q12) + 2*l2*m3*q12dot*q13dot*sin(2*q12) + 3*l1*m2*q11dot^2*sin(q11)*sin(q12)*sin(q13) + 2*l1*m3*q11dot^2*sin(q11)*sin(q12)*sin(q13)))/6;
        (l1*(3*g*m1*cos(q21) - 3*lambda2*sin(q21) - 6*lambda3*cos(q21) + 6*g*m2*cos(q21) + 2*g*m3*cos(q21) + 3*3^(1/2)*lambda1*sin(q21) - 3*l2*m2*q22dot^2*cos(q22)*sin(q21) - 2*l2*m3*q22dot^2*cos(q22)*sin(q21) + 3*l2*m2*q22dot^2*cos(q21)*cos(q23)*sin(q22) + 3*l2*m2*q23dot^2*cos(q21)*cos(q23)*sin(q22) + 2*l2*m3*q22dot^2*cos(q21)*cos(q23)*sin(q22) + 2*l2*m3*q23dot^2*cos(q21)*cos(q23)*sin(q22) + 6*l2*m2*q22dot*q23dot*cos(q21)*cos(q22)*sin(q23) + 4*l2*m3*q22dot*q23dot*cos(q21)*cos(q22)*sin(q23)))/6;
        (l2*(3*3^(1/2)*lambda1*sin(q22) - 3*lambda2*sin(q22) - 6*lambda3*cos(q22)*cos(q23) + 3*lambda1*cos(q22)*sin(q23) + l2*m2*q23dot^2*sin(2*q22) + l2*m3*q23dot^2*sin(2*q22) + 3*g*m2*cos(q22)*cos(q23) + 2*g*m3*cos(q22)*cos(q23) + 3*3^(1/2)*lambda2*cos(q22)*sin(q23) - 3*l1*m2*q21dot^2*cos(q21)*sin(q22) - 2*l1*m3*q21dot^2*cos(q21)*sin(q22) + 3*l1*m2*q21dot^2*cos(q22)*cos(q23)*sin(q21) + 2*l1*m3*q21dot^2*cos(q22)*cos(q23)*sin(q21)))/6;
        -(l2*(3*g*m2*sin(q22)*sin(q23) - 6*lambda3*sin(q22)*sin(q23) - 3*3^(1/2)*lambda2*cos(q23)*sin(q22) - 3*lambda1*cos(q23)*sin(q22) + 2*g*m3*sin(q22)*sin(q23) + 2*l2*m2*q22dot*q23dot*sin(2*q22) + 2*l2*m3*q22dot*q23dot*sin(2*q22) + 3*l1*m2*q21dot^2*sin(q21)*sin(q22)*sin(q23) + 2*l1*m3*q21dot^2*sin(q21)*sin(q22)*sin(q23)))/6;
        (l1*(3*g*m1*cos(q31) - 3*lambda5*sin(q31) - 6*lambda6*cos(q31) + 6*g*m2*cos(q31) + 2*g*m3*cos(q31) - 3*3^(1/2)*lambda4*sin(q31) - 3*l2*m2*q32dot^2*cos(q32)*sin(q31) - 2*l2*m3*q32dot^2*cos(q32)*sin(q31) + 3*l2*m2*q32dot^2*cos(q31)*cos(q33)*sin(q32) + 3*l2*m2*q33dot^2*cos(q31)*cos(q33)*sin(q32) + 2*l2*m3*q32dot^2*cos(q31)*cos(q33)*sin(q32) + 2*l2*m3*q33dot^2*cos(q31)*cos(q33)*sin(q32) + 6*l2*m2*q32dot*q33dot*cos(q31)*cos(q32)*sin(q33) + 4*l2*m3*q32dot*q33dot*cos(q31)*cos(q32)*sin(q33)))/6;
        (l2*(3*lambda4*cos(q32)*sin(q33) - 3*3^(1/2)*lambda4*sin(q32) - 6*lambda6*cos(q32)*cos(q33) - 3*lambda5*sin(q32) + l2*m2*q33dot^2*sin(2*q32) + l2*m3*q33dot^2*sin(2*q32) + 3*g*m2*cos(q32)*cos(q33) + 2*g*m3*cos(q32)*cos(q33) - 3*3^(1/2)*lambda5*cos(q32)*sin(q33) - 3*l1*m2*q31dot^2*cos(q31)*sin(q32) - 2*l1*m3*q31dot^2*cos(q31)*sin(q32) + 3*l1*m2*q31dot^2*cos(q32)*cos(q33)*sin(q31) + 2*l1*m3*q31dot^2*cos(q32)*cos(q33)*sin(q31)))/6;
        -(l2*(3*3^(1/2)*lambda5*cos(q33)*sin(q32) - 6*lambda6*sin(q32)*sin(q33) - 3*lambda4*cos(q33)*sin(q32) + 3*g*m2*sin(q32)*sin(q33) + 2*g*m3*sin(q32)*sin(q33) + 2*l2*m2*q32dot*q33dot*sin(2*q32) + 2*l2*m3*q32dot*q33dot*sin(2*q32) + 3*l1*m2*q31dot^2*sin(q31)*sin(q32)*sin(q33) + 2*l1*m3*q31dot^2*sin(q31)*sin(q32)*sin(q33)))/6
    ];

    h = [
        l2*sin(q12)*sin(q13) - (3^(1/2)*(r_base - r_platform + l1*cos(q21) + l2*cos(q22)))/2 + (l2*sin(q22)*sin(q23))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q21))/2 + (l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2
        l1*sin(q11) - l1*sin(q21) + l2*cos(q13)*sin(q12) - l2*cos(q23)*sin(q22)
        (3^(1/2)*(r_base - r_platform + l1*cos(q31) + l2*cos(q32)))/2 + l2*sin(q12)*sin(q13) + (l2*sin(q32)*sin(q33))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q31))/2 + (l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2
        l1*sin(q11) - l1*sin(q31) + l2*cos(q13)*sin(q12) - l2*cos(q33)*sin(q32)
    ];

    H1 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), (3^(1/2)*l1*sin(q21))/2, (3^(1/2)*l2*sin(q22))/2 + (l2*cos(q22)*sin(q23))/2, (l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H2 = [-l1*sin(q11), -l2*sin(q12), 0, -(l1*sin(q21))/2, (3^(1/2)*l2*cos(q22)*sin(q23))/2 - (l2*sin(q22))/2, (3^(1/2)*l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H3 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), -l1*cos(q21), -l2*cos(q22)*cos(q23), l2*sin(q22)*sin(q23), 0, 0, 0]';
    H4 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), 0, 0, 0, -(3^(1/2)*l1*sin(q31))/2, (l2*cos(q32)*sin(q33))/2 - (3^(1/2)*l2*sin(q32))/2, (l2*cos(q33)*sin(q32))/2]';
    H5 = [-l1*sin(q11), -l2*sin(q12), 0, 0, 0, 0, -(l1*sin(q31))/2, - (l2*sin(q32))/2 - (3^(1/2)*l2*cos(q32)*sin(q33))/2, -(3^(1/2)*l2*cos(q33)*sin(q32))/2]';
    H6 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), 0, 0, 0, -l1*cos(q31), -l2*cos(q32)*cos(q33), l2*sin(q32)*sin(q33)]';
    H = [H1; H2; H3; H4; H5; H6];

    H1dot = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), (3^(1/2)*l1*q21dot*cos(q21))/2, (l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*sin(q22)*sin(q23))/2 - (3^(1/2)*l2*cos(q22))/2), (l2*q22dot*cos(q22)*cos(q23))/2 - (l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    H2dot = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, -(l1*q21dot*cos(q21))/2, (3^(1/2)*l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2), (3^(1/2)*l2*q22dot*cos(q22)*cos(q23))/2 - (3^(1/2)*l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    H3dot = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), l1*q21dot*sin(q21), l2*q22dot*cos(q23)*sin(q22) + l2*q23dot*cos(q22)*sin(q23), l2*q22dot*cos(q22)*sin(q23) + l2*q23dot*cos(q23)*sin(q22), 0, 0, 0]';
    H4dot = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), 0, 0, 0, -(3^(1/2)*l1*q31dot*cos(q31))/2, (l2*q33dot*cos(q32)*cos(q33))/2 - q32dot*((l2*sin(q32)*sin(q33))/2 + (3^(1/2)*l2*cos(q32))/2), (l2*q32dot*cos(q32)*cos(q33))/2 - (l2*q33dot*sin(q32)*sin(q33))/2]';
    H5dot = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, 0, 0, 0, -(l1*q31dot*cos(q31))/2, - q32dot*((l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2) - (3^(1/2)*l2*q33dot*cos(q32)*cos(q33))/2, (3^(1/2)*l2*q33dot*sin(q32)*sin(q33))/2 - (3^(1/2)*l2*q32dot*cos(q32)*cos(q33))/2]';
    H6dot = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), 0, 0, 0, l1*q31dot*sin(q31), l2*q32dot*cos(q33)*sin(q32) + l2*q33dot*cos(q32)*sin(q33), l2*q32dot*cos(q32)*sin(q33) + l2*q33dot*cos(q33)*sin(q32)]';
    Hdot = [H1dot; H2dot; H3dot; H4dot; H5dot; H6dot];

    D = zeros(9);
    D[1] = p.d1;
    D[2] = p.d2;
    D[3] = p.d3;
    D[4] = p.d1;
    D[5] = p.d2;
    D[6] = p.d3;
    D[7] = p.d1;
    D[8] = p.d2;
    D[9] = p.d3;

    # invM_times_b = invM \ (b .- D .* x[10:18]); # invD_times_b = invD * b;
    invM_times_b = invM * (b .- D .* x[10:18]); # WTF
    z0 = h;
    z1 = H * x[10:18];
    z2 = Hdot * x[10:18] .+ H * invM_times_b;
    z = z0 + z1 + z2;

    out[1:9] = x[10:18] .- dx[1:9]; # first derivative residual
    out[10:18] = invM_times_b .- dx[10:18]; # second derivative residual
    out[19:24] = z; # index reduced holonomic constraint
end

"""
given the solution and the robot parameters, computes the lagrangian and hamiltonian for all time points an
"""
function delta_robot_energy(sol, p)

end

function gravity_compensated_model!(out, dx, x, p, t)
    # states of plant
    q11 = x[1];
    q12 = x[2];
    q13 = x[3];
    q21 = x[4];
    q22 = x[5];
    q23 = x[6];
    q31 = x[7];
    q32 = x[8];
    q33 = x[9];
    q11dot = x[10];
    q12dot = x[11];
    q13dot = x[12];
    q21dot = x[13];
    q22dot = x[14];
    q23dot = x[15];
    q31dot = x[16];
    q32dot = x[17];
    q33dot = x[18];
    lambda1 = x[19];
    lambda2 = x[20];
    lambda3 = x[21];
    lambda4 = x[22];
    lambda5 = x[23];
    lambda6 = x[24];

    # model params
    m1 = p.m1;
    m2 = p.m2;
    m3 = p.m3;
    l1 = p.l1;
    l2 = p.l2;
    g = p.g;
    r_base = p.rb;
    r_platform = p.rp

    # inertia matrix
    M1 = [(l1^ 2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6]';
    M2 = [(l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M3 = [-(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q12)^2 - 1))/3]';
    Mtop = [M1; M2; M3];

    M4 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6]';
    M5 = [(l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M6 = [-(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q22)^2 - 1))/3]';
    Mmid = [M4; M5; M6];

    M7 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6]';
    M8 = [(l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M9 = [-(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q32)^2 - 1))/3]';
    Mbot = [M7; M8; M9];

    invMtop = inv(Mtop);
    invMmid = inv(Mmid);
    invMbot = inv(Mbot);
    invM1 = [invMtop[1,:]; 0; 0; 0; 0; 0; 0]';
    invM2 = [invMtop[2,:]; 0; 0; 0; 0; 0; 0]';
    invM3 = [invMtop[3,:]; 0; 0; 0; 0; 0; 0]';
    invM4 = [0; 0; 0; invMmid[1,:]; 0; 0; 0]';
    invM5 = [0; 0; 0; invMmid[2,:]; 0; 0; 0]';
    invM6 = [0; 0; 0; invMmid[3,:]; 0; 0; 0]';
    invM7 = [0; 0; 0; 0; 0; 0; invMbot[1,:]]';
    invM8 = [0; 0; 0; 0; 0; 0; invMbot[2,:]]';
    invM9 = [0; 0; 0; 0; 0; 0; invMbot[3,:]]';
    invM = [invM1; invM2; invM3; invM4; invM5; invM6; invM7; invM8; invM9]

    C = [
        -(l1*l2*(3*m2 + 2*m3)*(q12dot^2*cos(q11)*cos(q13)*sin(q12) - q12dot^2*cos(q12)*sin(q11) + q13dot^2*cos(q11)*cos(q13)*sin(q12) + 2*q12dot*q13dot*cos(q11)*cos(q12)*sin(q13)))/6;
        -(l2*(l2*m2*q13dot^2*sin(2*q12) + l2*m3*q13dot^2*sin(2*q12) - 3*l1*m2*q11dot^2*cos(q11)*sin(q12) - 2*l1*m3*q11dot^2*cos(q11)*sin(q12) + 3*l1*m2*q11dot^2*cos(q12)*cos(q13)*sin(q11) + 2*l1*m3*q11dot^2*cos(q12)*cos(q13)*sin(q11)))/6;
        (l2*sin(q12)*(3*l1*m2*q11dot^2*sin(q11)*sin(q13) + 2*l1*m3*q11dot^2*sin(q11)*sin(q13) + 4*l2*m2*q12dot*q13dot*cos(q12) + 4*l2*m3*q12dot*q13dot*cos(q12)))/6;
        -(l1*l2*(3*m2 + 2*m3)*(q22dot^2*cos(q21)*cos(q23)*sin(q22) - q22dot^2*cos(q22)*sin(q21) + q23dot^2*cos(q21)*cos(q23)*sin(q22) + 2*q22dot*q23dot*cos(q21)*cos(q22)*sin(q23)))/6;
        -(l2*(l2*m2*q23dot^2*sin(2*q22) + l2*m3*q23dot^2*sin(2*q22) - 3*l1*m2*q21dot^2*cos(q21)*sin(q22) - 2*l1*m3*q21dot^2*cos(q21)*sin(q22) + 3*l1*m2*q21dot^2*cos(q22)*cos(q23)*sin(q21) + 2*l1*m3*q21dot^2*cos(q22)*cos(q23)*sin(q21)))/6;
        (l2*sin(q22)*(3*l1*m2*q21dot^2*sin(q21)*sin(q23) + 2*l1*m3*q21dot^2*sin(q21)*sin(q23) + 4*l2*m2*q22dot*q23dot*cos(q22) + 4*l2*m3*q22dot*q23dot*cos(q22)))/6;
        -(l1*l2*(3*m2 + 2*m3)*(q32dot^2*cos(q31)*cos(q33)*sin(q32) - q32dot^2*cos(q32)*sin(q31) + q33dot^2*cos(q31)*cos(q33)*sin(q32) + 2*q32dot*q33dot*cos(q31)*cos(q32)*sin(q33)))/6;
        -(l2*(l2*m2*q33dot^2*sin(2*q32) + l2*m3*q33dot^2*sin(2*q32) - 3*l1*m2*q31dot^2*cos(q31)*sin(q32) - 2*l1*m3*q31dot^2*cos(q31)*sin(q32) + 3*l1*m2*q31dot^2*cos(q32)*cos(q33)*sin(q31) + 2*l1*m3*q31dot^2*cos(q32)*cos(q33)*sin(q31)))/6;
        (l2*sin(q32)*(3*l1*m2*q31dot^2*sin(q31)*sin(q33) + 2*l1*m3*q31dot^2*sin(q31)*sin(q33) + 4*l2*m2*q32dot*q33dot*cos(q32) + 4*l2*m3*q32dot*q33dot*cos(q32)))/6
    ];

    h = [
        l2*sin(q12)*sin(q13) - (3^(1/2)*(r_base - r_platform + l1*cos(q21) + l2*cos(q22)))/2 + (l2*sin(q22)*sin(q23))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q21))/2 + (l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2
        l1*sin(q11) - l1*sin(q21) + l2*cos(q13)*sin(q12) - l2*cos(q23)*sin(q22)
        (3^(1/2)*(r_base - r_platform + l1*cos(q31) + l2*cos(q32)))/2 + l2*sin(q12)*sin(q13) + (l2*sin(q32)*sin(q33))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q31))/2 + (l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2
        l1*sin(q11) - l1*sin(q31) + l2*cos(q13)*sin(q12) - l2*cos(q33)*sin(q32)
    ];

    H1 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), (3^(1/2)*l1*sin(q21))/2, (3^(1/2)*l2*sin(q22))/2 + (l2*cos(q22)*sin(q23))/2, (l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H2 = [-l1*sin(q11), -l2*sin(q12), 0, -(l1*sin(q21))/2, (3^(1/2)*l2*cos(q22)*sin(q23))/2 - (l2*sin(q22))/2, (3^(1/2)*l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H3 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), -l1*cos(q21), -l2*cos(q22)*cos(q23), l2*sin(q22)*sin(q23), 0, 0, 0]';
    H4 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), 0, 0, 0, -(3^(1/2)*l1*sin(q31))/2, (l2*cos(q32)*sin(q33))/2 - (3^(1/2)*l2*sin(q32))/2, (l2*cos(q33)*sin(q32))/2]';
    H5 = [-l1*sin(q11), -l2*sin(q12), 0, 0, 0, 0, -(l1*sin(q31))/2, - (l2*sin(q32))/2 - (3^(1/2)*l2*cos(q32)*sin(q33))/2, -(3^(1/2)*l2*cos(q33)*sin(q32))/2]';
    H6 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), 0, 0, 0, -l1*cos(q31), -l2*cos(q32)*cos(q33), l2*sin(q32)*sin(q33)]';
    H = [H1; H2; H3; H4; H5; H6];

    H1dot = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), (3^(1/2)*l1*q21dot*cos(q21))/2, (l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*sin(q22)*sin(q23))/2 - (3^(1/2)*l2*cos(q22))/2), (l2*q22dot*cos(q22)*cos(q23))/2 - (l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    H2dot = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, -(l1*q21dot*cos(q21))/2, (3^(1/2)*l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2), (3^(1/2)*l2*q22dot*cos(q22)*cos(q23))/2 - (3^(1/2)*l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    H3dot = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), l1*q21dot*sin(q21), l2*q22dot*cos(q23)*sin(q22) + l2*q23dot*cos(q22)*sin(q23), l2*q22dot*cos(q22)*sin(q23) + l2*q23dot*cos(q23)*sin(q22), 0, 0, 0]';
    H4dot = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), 0, 0, 0, -(3^(1/2)*l1*q31dot*cos(q31))/2, (l2*q33dot*cos(q32)*cos(q33))/2 - q32dot*((l2*sin(q32)*sin(q33))/2 + (3^(1/2)*l2*cos(q32))/2), (l2*q32dot*cos(q32)*cos(q33))/2 - (l2*q33dot*sin(q32)*sin(q33))/2]';
    H5dot = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, 0, 0, 0, -(l1*q31dot*cos(q31))/2, - q32dot*((l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2) - (3^(1/2)*l2*q33dot*cos(q32)*cos(q33))/2, (3^(1/2)*l2*q33dot*sin(q32)*sin(q33))/2 - (3^(1/2)*l2*q32dot*cos(q32)*cos(q33))/2]';
    H6dot = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), 0, 0, 0, l1*q31dot*sin(q31), l2*q32dot*cos(q33)*sin(q32) + l2*q33dot*cos(q32)*sin(q33), l2*q32dot*cos(q32)*sin(q33) + l2*q33dot*cos(q33)*sin(q32)]';
    Hdot = [H1dot; H2dot; H3dot; H4dot; H5dot; H6dot];

    D = zeros(9);
    D[1] = p.d1;
    D[2] = p.d2;
    D[3] = p.d3;
    D[4] = p.d1;
    D[5] = p.d2;
    D[6] = p.d3;
    D[7] = p.d1;
    D[8] = p.d2;
    D[9] = p.d3;

    invM_times_b = invM * (-C .- D .* x[10:18]); #.- D .* x[10:18]); # invD_times_b = invD * b;
    z0 = h;
    z1 = H * x[10:18];
    z2 = Hdot * x[10:18] .+ H * invM_times_b;
    z = z0 + z1 + z2;

    out[1:9] = x[10:18] .- dx[1:9]; # first derivative residual
    out[10:18] = invM_times_b .- dx[10:18]; # second derivative residual
    out[19:24] = z; # index reduced holonomic constraint
end

function holonomic_constraint(q, p)
    q11 = q[1]
    q12 = q[2]
    q13 = q[3]
    q21 = q[4]
    q22 = q[5]
    q23 = q[6]
    q31 = q[7]
    q32 = q[8]
    q33 = q[9]
    m1 = p.m1;
    m2 = p.m2;
    m3 = p.m3;
    l1 = p.l1;
    l2 = p.l2;
    g = p.g;
    r_base = p.rb;
    r_platform = p.rp
    h = [
        l2*sin(q12)*sin(q13) - (3^(1/2)*(r_base - r_platform + l1*cos(q21) + l2*cos(q22)))/2 + (l2*sin(q22)*sin(q23))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q21))/2 + (l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2
        l1*sin(q11) - l1*sin(q21) + l2*cos(q13)*sin(q12) - l2*cos(q23)*sin(q22)
        (3^(1/2)*(r_base - r_platform + l1*cos(q31) + l2*cos(q32)))/2 + l2*sin(q12)*sin(q13) + (l2*sin(q32)*sin(q33))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q31))/2 + (l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2
        l1*sin(q11) - l1*sin(q31) + l2*cos(q13)*sin(q12) - l2*cos(q33)*sin(q32)
    ];
end

function L_delta_robot_base_pdff!(out, dx, x, p, t)
    # states of plant
    q11 = x[1];
    q12 = x[2];
    q13 = x[3];
    q21 = x[4];
    q22 = x[5];
    q23 = x[6];
    q31 = x[7];
    q32 = x[8];
    q33 = x[9];
    q11dot = x[10];
    q12dot = x[11];
    q13dot = x[12];
    q21dot = x[13];
    q22dot = x[14];
    q23dot = x[15];
    q31dot = x[16];
    q32dot = x[17];
    q33dot = x[18];
    lambda1 = x[19];
    lambda2 = x[20];
    lambda3 = x[21];
    lambda4 = x[22];
    lambda5 = x[23];
    lambda6 = x[24];

    # model params
    m1 = p.m1;
    m2 = p.m2;
    m3 = p.m3;
    l1 = p.l1;
    l2 = p.l2;
    g = p.g;
    r_base = p.rb;
    r_platform = p.rp

    # inertia matrix
    M1 = [(l1^ 2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6]';
    M2 = [(l1*l2*(sin(q11)*sin(q12) + cos(q11)*cos(q12)*cos(q13))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M3 = [-(l1*l2*cos(q11)*sin(q12)*sin(q13)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q12)^2 - 1))/3]';
    Mtop = [M1; M2; M3];

    M4 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6]';
    M5 = [(l1*l2*(sin(q21)*sin(q22) + cos(q21)*cos(q22)*cos(q23))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M6 = [-(l1*l2*cos(q21)*sin(q22)*sin(q23)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q22)^2 - 1))/3]';
    Mmid = [M4; M5; M6];

    M7 = [(l1^2*(m1 + 3*m2 + m3))/3, (l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, -(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6]';
    M8 = [(l1*l2*(sin(q31)*sin(q32) + cos(q31)*cos(q32)*cos(q33))*(3*m2 + 2*m3))/6, (l2^2*(m2 + m3))/3, 0]';
    M9 = [-(l1*l2*cos(q31)*sin(q32)*sin(q33)*(3*m2 + 2*m3))/6, 0, -(l2^2*(m2 + m3)*(cos(q32)^2 - 1))/3]';
    Mbot = [M7; M8; M9];

    invMtop = inv(Mtop);
    invMmid = inv(Mmid);
    invMbot = inv(Mbot);
    invM1 = [invMtop[1,:]; 0; 0; 0; 0; 0; 0]';
    invM2 = [invMtop[2,:]; 0; 0; 0; 0; 0; 0]';
    invM3 = [invMtop[3,:]; 0; 0; 0; 0; 0; 0]';
    invM4 = [0; 0; 0; invMmid[1,:]; 0; 0; 0]';
    invM5 = [0; 0; 0; invMmid[2,:]; 0; 0; 0]';
    invM6 = [0; 0; 0; invMmid[3,:]; 0; 0; 0]';
    invM7 = [0; 0; 0; 0; 0; 0; invMbot[1,:]]';
    invM8 = [0; 0; 0; 0; 0; 0; invMbot[2,:]]';
    invM9 = [0; 0; 0; 0; 0; 0; invMbot[3,:]]';
    invM = [invM1; invM2; invM3; invM4; invM5; invM6; invM7; invM8; invM9]

    b = [
        (l1*(6*lambda3*cos(q11) + 6*lambda6*cos(q11) - 6*lambda2*sin(q11) - 6*lambda5*sin(q11) + 3*g*m1*cos(q11) + 6*g*m2*cos(q11) + 2*g*m3*cos(q11) - 3*l2*m2*q12dot^2*cos(q12)*sin(q11) - 2*l2*m3*q12dot^2*cos(q12)*sin(q11) + 3*l2*m2*q12dot^2*cos(q11)*cos(q13)*sin(q12) + 3*l2*m2*q13dot^2*cos(q11)*cos(q13)*sin(q12) + 2*l2*m3*q12dot^2*cos(q11)*cos(q13)*sin(q12) + 2*l2*m3*q13dot^2*cos(q11)*cos(q13)*sin(q12) + 6*l2*m2*q12dot*q13dot*cos(q11)*cos(q12)*sin(q13) + 4*l2*m3*q12dot*q13dot*cos(q11)*cos(q12)*sin(q13)))/6;
        (l2*(6*lambda3*cos(q12)*cos(q13) - 6*lambda5*sin(q12) - 6*lambda2*sin(q12) + 6*lambda6*cos(q12)*cos(q13) + 6*lambda1*cos(q12)*sin(q13) + 6*lambda4*cos(q12)*sin(q13) + l2*m2*q13dot^2*sin(2*q12) + l2*m3*q13dot^2*sin(2*q12) + 3*g*m2*cos(q12)*cos(q13) + 2*g*m3*cos(q12)*cos(q13) - 3*l1*m2*q11dot^2*cos(q11)*sin(q12) - 2*l1*m3*q11dot^2*cos(q11)*sin(q12) + 3*l1*m2*q11dot^2*cos(q12)*cos(q13)*sin(q11) + 2*l1*m3*q11dot^2*cos(q12)*cos(q13)*sin(q11)))/6;
        -(l2*(6*lambda3*sin(q12)*sin(q13) - 6*lambda4*cos(q13)*sin(q12) - 6*lambda1*cos(q13)*sin(q12) + 6*lambda6*sin(q12)*sin(q13) + 3*g*m2*sin(q12)*sin(q13) + 2*g*m3*sin(q12)*sin(q13) + 2*l2*m2*q12dot*q13dot*sin(2*q12) + 2*l2*m3*q12dot*q13dot*sin(2*q12) + 3*l1*m2*q11dot^2*sin(q11)*sin(q12)*sin(q13) + 2*l1*m3*q11dot^2*sin(q11)*sin(q12)*sin(q13)))/6;
        (l1*(3*g*m1*cos(q21) - 3*lambda2*sin(q21) - 6*lambda3*cos(q21) + 6*g*m2*cos(q21) + 2*g*m3*cos(q21) + 3*3^(1/2)*lambda1*sin(q21) - 3*l2*m2*q22dot^2*cos(q22)*sin(q21) - 2*l2*m3*q22dot^2*cos(q22)*sin(q21) + 3*l2*m2*q22dot^2*cos(q21)*cos(q23)*sin(q22) + 3*l2*m2*q23dot^2*cos(q21)*cos(q23)*sin(q22) + 2*l2*m3*q22dot^2*cos(q21)*cos(q23)*sin(q22) + 2*l2*m3*q23dot^2*cos(q21)*cos(q23)*sin(q22) + 6*l2*m2*q22dot*q23dot*cos(q21)*cos(q22)*sin(q23) + 4*l2*m3*q22dot*q23dot*cos(q21)*cos(q22)*sin(q23)))/6;
        (l2*(3*3^(1/2)*lambda1*sin(q22) - 3*lambda2*sin(q22) - 6*lambda3*cos(q22)*cos(q23) + 3*lambda1*cos(q22)*sin(q23) + l2*m2*q23dot^2*sin(2*q22) + l2*m3*q23dot^2*sin(2*q22) + 3*g*m2*cos(q22)*cos(q23) + 2*g*m3*cos(q22)*cos(q23) + 3*3^(1/2)*lambda2*cos(q22)*sin(q23) - 3*l1*m2*q21dot^2*cos(q21)*sin(q22) - 2*l1*m3*q21dot^2*cos(q21)*sin(q22) + 3*l1*m2*q21dot^2*cos(q22)*cos(q23)*sin(q21) + 2*l1*m3*q21dot^2*cos(q22)*cos(q23)*sin(q21)))/6;
        -(l2*(3*g*m2*sin(q22)*sin(q23) - 6*lambda3*sin(q22)*sin(q23) - 3*3^(1/2)*lambda2*cos(q23)*sin(q22) - 3*lambda1*cos(q23)*sin(q22) + 2*g*m3*sin(q22)*sin(q23) + 2*l2*m2*q22dot*q23dot*sin(2*q22) + 2*l2*m3*q22dot*q23dot*sin(2*q22) + 3*l1*m2*q21dot^2*sin(q21)*sin(q22)*sin(q23) + 2*l1*m3*q21dot^2*sin(q21)*sin(q22)*sin(q23)))/6;
        (l1*(3*g*m1*cos(q31) - 3*lambda5*sin(q31) - 6*lambda6*cos(q31) + 6*g*m2*cos(q31) + 2*g*m3*cos(q31) - 3*3^(1/2)*lambda4*sin(q31) - 3*l2*m2*q32dot^2*cos(q32)*sin(q31) - 2*l2*m3*q32dot^2*cos(q32)*sin(q31) + 3*l2*m2*q32dot^2*cos(q31)*cos(q33)*sin(q32) + 3*l2*m2*q33dot^2*cos(q31)*cos(q33)*sin(q32) + 2*l2*m3*q32dot^2*cos(q31)*cos(q33)*sin(q32) + 2*l2*m3*q33dot^2*cos(q31)*cos(q33)*sin(q32) + 6*l2*m2*q32dot*q33dot*cos(q31)*cos(q32)*sin(q33) + 4*l2*m3*q32dot*q33dot*cos(q31)*cos(q32)*sin(q33)))/6;
        (l2*(3*lambda4*cos(q32)*sin(q33) - 3*3^(1/2)*lambda4*sin(q32) - 6*lambda6*cos(q32)*cos(q33) - 3*lambda5*sin(q32) + l2*m2*q33dot^2*sin(2*q32) + l2*m3*q33dot^2*sin(2*q32) + 3*g*m2*cos(q32)*cos(q33) + 2*g*m3*cos(q32)*cos(q33) - 3*3^(1/2)*lambda5*cos(q32)*sin(q33) - 3*l1*m2*q31dot^2*cos(q31)*sin(q32) - 2*l1*m3*q31dot^2*cos(q31)*sin(q32) + 3*l1*m2*q31dot^2*cos(q32)*cos(q33)*sin(q31) + 2*l1*m3*q31dot^2*cos(q32)*cos(q33)*sin(q31)))/6;
        -(l2*(3*3^(1/2)*lambda5*cos(q33)*sin(q32) - 6*lambda6*sin(q32)*sin(q33) - 3*lambda4*cos(q33)*sin(q32) + 3*g*m2*sin(q32)*sin(q33) + 2*g*m3*sin(q32)*sin(q33) + 2*l2*m2*q32dot*q33dot*sin(2*q32) + 2*l2*m3*q32dot*q33dot*sin(2*q32) + 3*l1*m2*q31dot^2*sin(q31)*sin(q32)*sin(q33) + 2*l1*m3*q31dot^2*sin(q31)*sin(q32)*sin(q33)))/6
    ];

    h = [
        l2*sin(q12)*sin(q13) - (3^(1/2)*(r_base - r_platform + l1*cos(q21) + l2*cos(q22)))/2 + (l2*sin(q22)*sin(q23))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q21))/2 + (l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2
        l1*sin(q11) - l1*sin(q21) + l2*cos(q13)*sin(q12) - l2*cos(q23)*sin(q22)
        (3^(1/2)*(r_base - r_platform + l1*cos(q31) + l2*cos(q32)))/2 + l2*sin(q12)*sin(q13) + (l2*sin(q32)*sin(q33))/2
        (3*r_base)/2 - (3*r_platform)/2 + l1*cos(q11) + l2*cos(q12) + (l1*cos(q31))/2 + (l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2
        l1*sin(q11) - l1*sin(q31) + l2*cos(q13)*sin(q12) - l2*cos(q33)*sin(q32)
    ];

    H1 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), (3^(1/2)*l1*sin(q21))/2, (3^(1/2)*l2*sin(q22))/2 + (l2*cos(q22)*sin(q23))/2, (l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H2 = [-l1*sin(q11), -l2*sin(q12), 0, -(l1*sin(q21))/2, (3^(1/2)*l2*cos(q22)*sin(q23))/2 - (l2*sin(q22))/2, (3^(1/2)*l2*cos(q23)*sin(q22))/2, 0, 0, 0]';
    H3 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), -l1*cos(q21), -l2*cos(q22)*cos(q23), l2*sin(q22)*sin(q23), 0, 0, 0]';
    H4 = [0, l2*cos(q12)*sin(q13), l2*cos(q13)*sin(q12), 0, 0, 0, -(3^(1/2)*l1*sin(q31))/2, (l2*cos(q32)*sin(q33))/2 - (3^(1/2)*l2*sin(q32))/2, (l2*cos(q33)*sin(q32))/2]';
    H5 = [-l1*sin(q11), -l2*sin(q12), 0, 0, 0, 0, -(l1*sin(q31))/2, - (l2*sin(q32))/2 - (3^(1/2)*l2*cos(q32)*sin(q33))/2, -(3^(1/2)*l2*cos(q33)*sin(q32))/2]';
    H6 = [l1*cos(q11), l2*cos(q12)*cos(q13), -l2*sin(q12)*sin(q13), 0, 0, 0, -l1*cos(q31), -l2*cos(q32)*cos(q33), l2*sin(q32)*sin(q33)]';
    H = [H1; H2; H3; H4; H5; H6];

    H1dot = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), (3^(1/2)*l1*q21dot*cos(q21))/2, (l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*sin(q22)*sin(q23))/2 - (3^(1/2)*l2*cos(q22))/2), (l2*q22dot*cos(q22)*cos(q23))/2 - (l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    H2dot = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, -(l1*q21dot*cos(q21))/2, (3^(1/2)*l2*q23dot*cos(q22)*cos(q23))/2 - q22dot*((l2*cos(q22))/2 + (3^(1/2)*l2*sin(q22)*sin(q23))/2), (3^(1/2)*l2*q22dot*cos(q22)*cos(q23))/2 - (3^(1/2)*l2*q23dot*sin(q22)*sin(q23))/2, 0, 0, 0]';
    H3dot = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), l1*q21dot*sin(q21), l2*q22dot*cos(q23)*sin(q22) + l2*q23dot*cos(q22)*sin(q23), l2*q22dot*cos(q22)*sin(q23) + l2*q23dot*cos(q23)*sin(q22), 0, 0, 0]';
    H4dot = [0, l2*q13dot*cos(q12)*cos(q13) - l2*q12dot*sin(q12)*sin(q13), l2*q12dot*cos(q12)*cos(q13) - l2*q13dot*sin(q12)*sin(q13), 0, 0, 0, -(3^(1/2)*l1*q31dot*cos(q31))/2, (l2*q33dot*cos(q32)*cos(q33))/2 - q32dot*((l2*sin(q32)*sin(q33))/2 + (3^(1/2)*l2*cos(q32))/2), (l2*q32dot*cos(q32)*cos(q33))/2 - (l2*q33dot*sin(q32)*sin(q33))/2]';
    H5dot = [-l1*q11dot*cos(q11), -l2*q12dot*cos(q12), 0, 0, 0, 0, -(l1*q31dot*cos(q31))/2, - q32dot*((l2*cos(q32))/2 - (3^(1/2)*l2*sin(q32)*sin(q33))/2) - (3^(1/2)*l2*q33dot*cos(q32)*cos(q33))/2, (3^(1/2)*l2*q33dot*sin(q32)*sin(q33))/2 - (3^(1/2)*l2*q32dot*cos(q32)*cos(q33))/2]';
    H6dot = [-l1*q11dot*sin(q11), - l2*q12dot*cos(q13)*sin(q12) - l2*q13dot*cos(q12)*sin(q13), - l2*q12dot*cos(q12)*sin(q13) - l2*q13dot*cos(q13)*sin(q12), 0, 0, 0, l1*q31dot*sin(q31), l2*q32dot*cos(q33)*sin(q32) + l2*q33dot*cos(q32)*sin(q33), l2*q32dot*cos(q32)*sin(q33) + l2*q33dot*cos(q33)*sin(q32)]';
    Hdot = [H1dot; H2dot; H3dot; H4dot; H5dot; H6dot];

    D = zeros(9);
    D[1] = p.d1;
    D[2] = p.d2;
    D[3] = p.d3;
    D[4] = p.d1;
    D[5] = p.d2;
    D[6] = p.d3;
    D[7] = p.d1;
    D[8] = p.d2;
    D[9] = p.d3;

    lambda = [lamdba1; lambda2; lambda3; lambda4; lambda5; lambda6];
    u = -b .+ D .* x[10:18] .+ H' * lambda;
    u = u .+ inv(invM) * 0.001 * [deg2rad(20)-q11; 0; 0; deg2rad(20)-q21; 0; 0; deg2rad(20)-q31; 0; 0];

    invM_times_b = invM * (b .- D .* x[10:18] .+ u);
    z0 = h;
    z1 = H * x[10:18];
    z2 = Hdot * x[10:18] .+ H * invM_times_b;
    z = z0 + z1 + z2;

    out[1:9] = x[10:18] .- dx[1:9]; # first derivative residual
    out[10:18] = invM_times_b .- dx[10:18]; # second derivative residual
    out[19:24] = z; # index reduced holonomic constraint
end

max_stars_repo_name: zborffs/Delta

max_issues_repo_name: zborffs/Delta

max_forks_repo_name: zborffs/Delta


text:
mutable struct ILU0Preconditioner{Tv, Ti} <: AbstractExtendablePreconditioner{Tv,Ti}
    extmatrix::ExtendableSparseMatrix{Tv,Ti}
    xdiag::Array{Tv,1}
    idiag::Array{Ti,1}
    pattern_timestamp::Float64
end

function ILU0Preconditioner(extmatrix::ExtendableSparseMatrix{Tv,Ti}) where {Tv,Ti}
    @assert size(extmatrix,1)==size(extmatrix,2)
    flush!(extmatrix)
    n=size(extmatrix,1)
    xdiag=Array{Tv,1}(undef,n)
    idiag=Array{Ti,1}(undef,n)
    precon=ILU0Preconditioner{Tv, Ti}(extmatrix,xdiag,idiag,0.0)
    update!(precon)
end

ILU0Preconditioner(cscmatrix::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}=ILU0Preconditioner(ExtendableSparseMatrix(cscmatrix))

function update!(precon::ILU0Preconditioner{Tv,Ti}) where {Tv,Ti}
    cscmatrix=precon.extmatrix.cscmatrix
    colptr=cscmatrix.colptr
    rowval=cscmatrix.rowval
    nzval=cscmatrix.nzval
    n=cscmatrix.n
    xdiag=precon.xdiag
    idiag=precon.idiag

    # Find main diagonal index and
    # copy main diagonal values
    if need_symbolic_update(precon)
        @inbounds for j=1:n
            @inbounds for k=colptr[j]:colptr[j+1]-1
                i=rowval[k]
                if i==j
                    idiag[j]=k
                    break
                end
            end
        end
        timestamp!(precon)
    end

    @inbounds for j=1:n
        xdiag[j]=one(Tv)/nzval[idiag[j]]
        @inbounds for k=idiag[j]+1:colptr[j+1]-1
            i=rowval[k]
            for l=colptr[i]:colptr[i+1]-1
                if rowval[l]==j
                    xdiag[i]-=nzval[l]*xdiag[j]*nzval[k]
                    break
                end
            end
        end
    end
    precon
end

function  LinearAlgebra.ldiv!(u::AbstractArray{T,1}, precon::ILU0Preconditioner, v::AbstractArray{T,1}) where T
    cscmatrix=precon.extmatrix.cscmatrix
    colptr=cscmatrix.colptr
    rowval=cscmatrix.rowval
    n=cscmatrix.n
    nzval=cscmatrix.nzval
    xdiag=precon.xdiag
    idiag=precon.idiag

    @inbounds for j=1:n
        x=zero(T)
        @inbounds for k=colptr[j]:idiag[j]-1
            x+=nzval[k]*u[rowval[k]]
        end
        u[j]=xdiag[j]*(v[j]-x)
    end

    @inbounds for j=n:-1:1
        x=zero(T)
        @inbounds for k=idiag[j]+1:colptr[j+1]-1
            x+=u[rowval[k]]*nzval[k]
        end
        u[j]-=x*xdiag[j]
    end
end

ext: jl


text:
# we define our own compatibilty layer here until Compat.jl#261 is merged
if isdefined(Base, :AsyncCondition)
    import Base.AsyncCondition
else
    type AsyncCondition
        cond::Condition
        handle::Ptr{Void}

        function AsyncCondition(func=nothing)
            this = new(Condition())
            # the callback is supposed to be called with this AsyncCondition
            # as the argument, so we need to create a wrapper callback
            if func == nothing
                function wrapfunc(data)
                    notify(this.cond)

                    nothing
                end
            else
                function wrapfunc(data)
                    notify(this.cond)
                    func(this)

                    nothing
                end
            end
            work = Base.SingleAsyncWork(wrapfunc)
            this.handle = work.handle

            this
        end
    end

    Base.wait(c::AsyncCondition) = wait(c.cond)
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""Represents a volume that is populated with the contents of a git repository. Git repo volumes do not support ownership management. Git repo volumes support SELinux relabeling.  DEPRECATED: GitRepo is deprecated. To provision a container with a git repo, mount an EmptyDir into an InitContainer that clones the repo using git, then mount the EmptyDir into the Pod's container.

    IoK8sApiCoreV1GitRepoVolumeSource(;
        directory=nothing,
        repository=nothing,
        revision=nothing,
    )

    - directory::String : Target directory name. Must not contain or start with '..'.  If '.' is supplied, the volume directory will be the git repository.  Otherwise, if specified, the volume will contain the git repository in the subdirectory with the given name.
    - repository::String : Repository URL
    - revision::String : Commit hash for the specified revision.
"""
mutable struct IoK8sApiCoreV1GitRepoVolumeSource <: SwaggerModel
    directory::Any # spec type: Union{ Nothing, String } # spec name: directory
    repository::Any # spec type: Union{ Nothing, String } # spec name: repository
    revision::Any # spec type: Union{ Nothing, String } # spec name: revision

    function IoK8sApiCoreV1GitRepoVolumeSource(;directory=nothing, repository=nothing, revision=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1GitRepoVolumeSource, Symbol("directory"), directory)
        setfield!(o, Symbol("directory"), directory)
        validate_property(IoK8sApiCoreV1GitRepoVolumeSource, Symbol("repository"), repository)
        setfield!(o, Symbol("repository"), repository)
        validate_property(IoK8sApiCoreV1GitRepoVolumeSource, Symbol("revision"), revision)
        setfield!(o, Symbol("revision"), revision)
        o
    end
end # type IoK8sApiCoreV1GitRepoVolumeSource

const _property_map_IoK8sApiCoreV1GitRepoVolumeSource = Dict{Symbol,Symbol}(Symbol("directory")=>Symbol("directory"), Symbol("repository")=>Symbol("repository"), Symbol("revision")=>Symbol("revision"))
const _property_types_IoK8sApiCoreV1GitRepoVolumeSource = Dict{Symbol,String}(Symbol("directory")=>"String", Symbol("repository")=>"String", Symbol("revision")=>"String")
Base.propertynames(::Type{ IoK8sApiCoreV1GitRepoVolumeSource }) = collect(keys(_property_map_IoK8sApiCoreV1GitRepoVolumeSource))
Swagger.property_type(::Type{ IoK8sApiCoreV1GitRepoVolumeSource }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiCoreV1GitRepoVolumeSource[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1GitRepoVolumeSource }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1GitRepoVolumeSource[property_name]

function check_required(o::IoK8sApiCoreV1GitRepoVolumeSource)
    (getproperty(o, Symbol("repository")) === nothing) && (return false)
    true
end

function validate_property(::Type{ IoK8sApiCoreV1GitRepoVolumeSource }, name::Symbol, val)
end


text:
using CSV, KernelDensity, Plots, DataFrames

"""
    get_experimental(filename, ν0_guess)

Get the experimental data from the filename specified and with the guess for the
larmor frequency (ν0_guess). Returns data of type Array{Float64,2} assuming good
data in the specified file.

# Examples
```julia-repl
julia> experimental = get_experimental("B2O3FastCool.txt", 32.239)
9992×2 Array{Float64,2}:
 30.99     8158.06
 30.9902   7660.22
 30.9905   7280.4
 30.9907   7220.96
 30.991    7524.46
 30.9912   8042.03
 30.9915   8545.74
 30.9917   8843.87
 30.992    8826.42
 30.9922   8468.01
 30.9924   7816.08
 30.9927   6973.43
 30.9929   6075.45
  ⋮
 33.4273  16466.1
 33.4276  15877.6
 33.4278  15079.3
 33.4281  14120.7
 33.4283  13077.5
 33.4285  12044.6
 33.4288  11128.3
 33.429   10433.0
 33.4293  10042.0
 33.4295   9993.03
 33.4298  10261.3
 33.43    10764.6
```
"""

function get_experimental(
    filename::String;
    ν0_guess::Float64 = 0.0,
    header::Bool = false,
    delim = ",",
    convert_ppm_to_mhz::Bool = ν0_guess != 0,
)
    experimental = CSV.read(filename, DataFrame, header = header, delim = delim)
    rename!(experimental, [:Frequency, :Intensity])
    if typeof(experimental[1, 1]) <: AbstractString
        experimental[!, 1] = parse.(Float64, experimental[:, 1])
    end
    if convert_ppm_to_mhz
        experimental[!, 1] = experimental[:, 1] .* (ν0_guess / (10^6)) .+ ν0_guess
    end
    sort!(experimental, [:Frequency])

    return Matrix(experimental)
end

function get_data(filename::String)
    return joinpath(dirname(pathof(SpectraFit)), "..", "data", filename)
end

"""
    generate_theoretical_spectrum(experimental, Quadrupolar)

Given the experimental data (experimental) and a set of corresponding
nmr parameters, constructs a powder pattern (spectrum) by interpolating the
kernel density estimate of the distribution produced by the nmr parameters

# Examples
```julia-repl
julia> generate_theoretical_spectrum(
julia>     get_experimental("B2O3FastCool.txt", 32.239),
julia>     Quadrupolar([5.5, 0.1, 0.12, 0.03, 1.0])
9992-element Array{Float64,1}:
 20431.55828260406
 20423.94683368574
 20415.98848956408
 20408.319911200455
 20400.302057265526
 20392.58886476615
 20384.5664157466
 20376.877421554564
 20369.200721475398
 20361.217235488795
 20353.565635884566
 20345.621287321344
 20338.045108945793
     ⋮
 22608.80291238082
 22607.488079524057
 22606.341148346466
 22605.26654097029
 22604.350227759416
 22603.546870237795
 22602.831135390592
 22602.261303660813
 22601.78989055206
 22601.454608723958
 22601.22751832443
 22601.126786321012
```
"""
function generate_theoretical_spectrum(
    experimental::Array{Float64,2},
    Quadrupolar::Quadrupolar;
    I::Int64 = 3,
    transitions = 1:(2*I),
)
    ν0 = get_ν0(experimental, get_experimental_ecdf(experimental))
    powder_pattern = estimate_powder_pattern(Quadrupolar, 1_000_000, ν0, I,
        transitions = transitions)
    k = kde(powder_pattern)
    x = experimental[:, 1]
    ik = InterpKDE(k)
    theoretical = pdf(ik, x)
    return (mean(experimental[:, 2]) / mean(theoretical)) .* theoretical
end

function generate_theoretical_spectrum(
    experimental::Array{Float64,2},
    ChemicalShift::ChemicalShift,
)
    powder_pattern = estimate_powder_pattern(ChemicalShift, 1_000_000)
    k = kde(powder_pattern)
    x = experimental[:, 1]
    ik = InterpKDE(k)
    theoretical = pdf(ik, x)
    return (mean(experimental[:, 2]) / mean(theoretical)) .* theoretical
end

function plot_experimental(
    experimental::Array{Float64,2};
    relative_ν = false,
    unit = "MHz",
)
    to_plot = relative_ν ?
        experimental[:, 1] .- get_ν(experimental) :
        experimental[:, 1]
    x_label = relative_ν ? "Δν ($(unit))" : "ν ($(unit))"
    plot(
        to_plot,
        experimental[:, 2],
        label = "experimental",
        xlabel = x_label,
        ylabel = "Intensity",
    )
end

function plot_theoretical(
    experimental::Array{Float64,2},
    params::Quadrupolar;
    relative_ν = false,
)
    to_plot = relative_ν ?
        experimental[:, 1] .- get_ν(experimental) :
        experimental[:, 1]
    x_label = relative_ν ? "Δν (MHz)" : "ν (MHz)"
    theoretical = generate_theoretical_spectrum(experimental, params)
    plot(
        to_plot,
        theoretical,
        width = 2,
        label = "theoretical",
        xlabel = x_label,
        ylabel = "Intensity",
    )
end

function plot_theoretical(
    experimental::Array{Float64,2},
    params::ChemicalShift;
    relative_ν = false,
    unit = "MHz",
)
    to_plot = relative_ν ?
        experimental[:, 1] .- get_ν(experimental) :
        experimental[:, 1]
    x_label = relative_ν ? "Δν ($(unit))" : "ν ($(unit))"
    theoretical = generate_theoretical_spectrum(experimental, params)
    plot(
        to_plot,
        theoretical,
        width = 2,
        label = "theoretical",
        xlabel = x_label,
        ylabel = "Intensity",
    )
end

"""
    compare_theoreticals(experimental, old_Quadrupolar, new_Quadrupolar)

Plots the experimental data along with two theoretical powder patterns, for the
sake of comparison

"""
function compare_theoreticals(
    experimental::Array{Float64,2},
    old_Quadrupolar::Quadrupolar,
    new_Quadrupolar::Quadrupolar;
    I = 3,
    transitions::UnitRange{Int64} = 1:(2*I),
)
    plot(experimental[:, 1], experimental[:, 2], label = "experimental")
    theoretical = generate_theoretical_spectrum(experimental, old_Quadrupolar,
        transitions = transitions)
    plot!(experimental[:, 1], theoretical, width = 2, label = "old theoretical")
    theoretical = generate_theoretical_spectrum(experimental, new_Quadrupolar,
        transitions = transitions)
    plot!(
        experimental[:, 1],
        theoretical,
        width = 2,
        label = "new theoretical",
        title = "Theoretical vs. Experimental",
        xlabel = "Frequency (MHz)",
        ylabel = "Intensity",
    )
end

function plot_chemical_shift(
    experimental::Array{Float64,2},
    params::ChemicalShift
)
    plot(experimental[:, 1], experimental[:, 2], label = "experimental")
    theoretical = generate_theoretical_spectrum(experimental, params)
    plot!(
        experimental[:, 1],
        theoretical,
        width = 2,
        label = "theoretical",
        title = "Theoretical vs. Experimental",
        xlabel = "Frequency (ppm)",
        ylabel = "Intensity",
    )
end

function plot_chemical_shift(
    experimental::Array{Float64,2},
    params::Array{Float64}
)
    return plot_chemical_shift(experimental,
        ChemicalShift(transform_params(push!(params, 1.0), ChemicalShift)))
end

"""
    get_experimental_ecdf(experimental)

Returns the ecdf of the experimental data over the range of the experimental
data

"""
function get_experimental_ecdf(experimental::Array{Float64,2})
    return cumsum(experimental[:, 2]) ./ sum(experimental[:, 2])
end

"""
    get_ν0(experimental, experimental_ecdf)

Calculates the Larmor frequency by finding the approximate mean of the
distribution. Should be passed into the forward model to get theoretical data
that lines up with the experimental

"""
function get_ν0(experimental::Array{Float64,2}, experimental_ecdf)
    riemann_sum = 0
    for i = 2:length(experimental_ecdf)
        riemann_sum += (experimental_ecdf[i]) *
                       (experimental[i, 1] - experimental[i-1, 1])
    end
    return experimental[end, 1] - riemann_sum
end

function get_ν0(experimental::Array{Float64,2})
    return get_ν0(experimental, get_experimental_ecdf(experimental))
end


text:
# This file was generated, do not modify it. # hide
r3(x) = round(x, sigdigits=3)
r3(pi)

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
n = rand(1:10)
@test matrixdepot("cauchy", n) == matrixdepot("cauchy", Float64[1:n;])
x = ones(Float64, n)
y = x + 2
A = zeros(Float64, n, n)
fill!(A, 0.25)
@test_approx_eq matrixdepot("cauchy", x, y) A
println("'cauchy' passed test...")


text:
using Test

@testset "model.jl" begin
    include("influence_diagram.jl")
    include("random.jl")
    include("decision_model.jl")
end


text:
mutable struct DMHTune <: SamplerTune
    logf::Union{Function,Missing}
    #get_model::Function
    pseudolog::Union{Function,Missing} # AuxLog
    condlike::Union{Function,Missing}
    datakeys::Vector{Symbol}
    m::Int64
    scale::Float64

    DMHTune() = new()

    function DMHTune(f::Union{Function,Missing}, ps, cl, m::Int64, s::Float64)
        new(f, ps, cl, Vector{Symbol}(), m, s)
    end
    function DMHTune(
        f::Union{Function,Missing},
        ps,
        cl,
        dk::Vector{Symbol},
        m::Int64,
        s::Float64,
    )
        new(f, ps, cl, dk, m, s)
    end
end

const DMHVariate = Sampler{DMHTune,T} where {T}

DMHTune(x::Vector, f, ps, cl, dk, m, s; args...) = DMHTune(f, ps, cl, dk, m, s; args...)

#################### Sampler Constructor ####################

function DMH(
    params::ElementOrVector{Symbol},
    m::Int64,
    scale::Float64 = 1.0,
    transform::Bool = true;
    args...,
)

    tune = DMHTune(logpdf!, pseudologpdf!, conditional_likelihood!, m, scale)

    Sampler(params, tune, Symbol[], transform)
end

#  sample!(v::DMHVariate, model) = DMH_sample!(v, v.tune, model)

function sample!(v::DMHVariate{T}, lpdf::Function; model::Model, kwargs...) where {T}#tune::DMHTune, model::Model)
    # 1. propose theta prime
    # 2. Generate the auxiliary variable using theta prime

    tune = v.tune
    #model = get_model(0)
    #lpdf(x) = logpdf!(model, x, v.params, v.targets, v.transform)
    targets = keys(model, :target, v.params)
    data = unlist(model[targets[1]])
    data = reshape(data, size(getindex(model, targets[1])))
    nfeatures, nlang = size(data)
    #xobs = Array{Any}(undef, nfeatures)
    #nobs = Array{Float64}(undef, nfeatures)

    # Separate missing from observed indices
    #for f in 1:nfeatures
    #	xmis = findall(x -> x .== -10, data[f,:])
    #	xobs[f,:] = findall(x -> x .≠ -10, data[f,:])
    #	nobs[f] = length(xobs[f,:])
    #end

    @assert v.tune.m >= nlang

    # store old value for possible future reference
    θ = deepcopy(v.value)
    # this way of generating theta_prime from the current values of theta
    # takes care of the transition probability from theta_prime to theta and vice versa
    # the values equal and will cancel out.
    θ_prime = θ + v.tune.scale .* rand(Normal(0.0, 1.0), length(v))

    #println("real")
    # calculate logpdf values
    lf_xt = 0.0#lpdf(v.value)

    #println("θ_prime $θ_prime")
    lf_xtp = 0.0#lpdf(θ_prime)

    # prematurely assign, to allow computations to go through properly
    v[:] = θ_prime

    # Sample new pseudo observations based on the θ_prime
    y = inner_sampler(v, deepcopy(data), v.params, targets, model) #xobs

    # calculate logpdfs based on pseudo observations
    lf_ytp = tune.pseudolog(model, v.value, y, v.params, targets)
    lf_yt = tune.pseudolog(model, θ, y, v.params, targets)

    # calculate acceptance probability (proposal distribution?)
    r = exp((lf_yt + lf_xtp) - (lf_xt + lf_ytp))

    # RWM acceptance logic is a bit reversed here. Above the new value is
    # prematurely assigned to allow computations in the inner sampler to go through.
    # If the sample is accepted nothing needs to be done anymore, otherwise the
    # old value will be reassigned.

    if rand() > r
        # sample is rejected, so use the old value.
        v[:] = θ
    end
    v
end

function inner_sampler(
    v::DMHVariate,
    X::Array{N,2},
    params,
    targets,
    model,
)::Array{N,2} where {N<:Real}
    nfeatures, nlangs = size(X)
    counter = zero(Int64)

    while true
        #random_language_idx = shuffle(1:nlangs)
        random_feature_idx = shuffle(1:nfeatures)
        @inbounds for i = 1:nlangs#random_language_idx
            @inbounds for f in random_feature_idx

                probs = v.tune.condlike(model, X, params, targets, i, f)
                new_x = sample(1:length(probs), Weights(probs))
                X[f, i] = new_x
                #	end

                counter += 1
                if counter > v.tune.m
                    return X
                end
            end
        end
    end
end

max_stars_repo_name: erathorn/MCPhylo

max_issues_repo_name: erathorn/MCPhylo

max_forks_repo_name: erathorn/MCPhylo


text:
##############################################################################
##
## simulated_annealing()
##
## Arguments:
## * cost: Function from states to the real numbers. Often called an energy
##         function, but this algorithm works for both positive and negative
##         costs.
## * s0: The initial state of the system.
## * neighbor: Function from states to states. Produces what the Metropolis
##             algorithm would call a proposal.
## * temperature: Function specifying the temperature at time i.
## * iterations: How many iterations of the algorithm should be run? This is
##               the only termination condition.
## * keep_best: Do we return the best state visited or the last state visited?
##              (Should default to true.)
## * show_trace: Do we show a trace of the system's evolution?
##
##############################################################################

##############################################################################
##
## Off-the-shelf cooling schedules
##
##############################################################################

# Theoretically, SA will converge if the temperature decreases
# according inversely proportional to the current time.
#
# For this to work, an unknown constant must be used.
# In practice, we use 1 instead of this unknown constant.
log_temperature(t::Real) = 1.0 / log(t)

# If the temperature is held constant and the cost function is -log(p),
# SA reduces to the Metropolis algorithm for sampling from a distribution.
constant_temperature(t::Real) = 1.0

# Default neighbor takes Gaussian jumps
function default_neighbor!(x::Vector, storage::Vector)
    for i in 1:length(x)
        storage[i] = x[i] + randn()
    end
end

# NB: The neighbor function is mutating now
function simulated_annealing(cost::Function,
                             s0::Vector;
                             neighbor!::Function = default_neighbor!,
                             temperature::Function = log_temperature,
                             keep_best::Bool = true,
                             tolerance::Real = 1e-8,
                             iterations::Integer = 100_000,
                             store_trace::Bool = false,
                             show_trace::Bool = false)

    # Maintain a trace of the optimization algo's state
    tr = OptimizationTrace()

    # Set our current state to the specified intial state
    s = copy(s0)
    s_n = copy(s0)
    y = cost(s)

    # Track how many function calls we've made
    f_calls = 1

    # Set the best state we've seen to the intial state
    best_s = copy(s0)
    best_y = y

    # Record the number of iterations we perform
    iteration = 0

    # Update our trace information
    os = OptimizationState(copy(s), y, iteration)
    if store_trace
        push!(tr, os)
    end
    if show_trace
        println(os)
    end

    # We always perform a fixed number of iterations
    while iteration < iterations
        # Update the iteration counter
        iteration += 1

        # Call temperature to find the proper temperature at time i
        t = temperature(iteration)

        # Call neighbor to randomly generate a neighbor of our current state
        neighbor!(s, s_n)

        # Evaluate the cost function on our current state and its neighbor
        y_n = cost(s_n)
        f_calls += 1

        if y_n <= y
            # If the proposed new state is superior, we always move to it
            copy!(s, s_n)
            y = y_n

            # If the new state is the best state we have seen,
            #  keep a record of it
            if y_n < best_y
                best_y = y_n
                copy!(best_s, s_n)
            end
        else
            # If the proposed new state is inferior, we move to it with
            #  probability p
            p = exp(-(y_n - y)/t)
            if rand() <= p
                copy!(s, s_n)
                y = y_n
            end
        end

        # Print out the state of the system
        if store_trace || show_trace
            os = OptimizationState(copy(s), y, iteration)
            if store_trace
                push!(tr, os)
            end
            if show_trace
                println(os)
            end
        end
    end

    # If specified by the user, we return the best state we've seen
    # Otherwise, we return the last state we've seen
    if keep_best
        OptimizationResults("Simulated Annealing",
                            s0,
                            best_s,
                            best_y,
                            iterations,
                            false,
                            tr,
                            f_calls,
                            0)
    else
        OptimizationResults("Simulated Annealing",
                            s0,
                            s,
                            y,
                            iterations,
                            false,
                            tr,
                            f_calls,
                            0)
    end
end


text:
using ObjectiveFunctions
using Base.Test

@testset "RegObj" begin
    let nin=2, nout=3
        t = Affine(nin, nout)
        l = L2DistLoss()
        λ = 1e-3
        p = scaled(L2Penalty(), λ)
        obj = objective(t, l, p)
        @show obj typeof(obj)

        @test output_value(t) == input_value(obj.loss)
        @test input_length(obj.loss) == nout
        @test output_length(obj.loss) == 1

        input = rand(nin)
        target = rand(nout)
        w, b = t.params.views

        # test the forward pass
        transform!(obj, target, input)
        @test input_value(obj) == input
        @test output_value(t) ≈ w * input + b
        total_loss = sum(value(l, target[i], output_value(t)[i]) for i=1:nout)
        total_penalty = 0.5λ * sum(θᵢ^2 for θᵢ in params(t))
        @show total_loss total_penalty
        @test output_value(obj)[1] ≈ total_loss + total_penalty

        # test the backward pass
        grad!(obj)
        x = input_value(obj.loss)
        ∇ = input_grad(obj.loss)
        @show ∇
        for i=1:nout
            @test ∇[i] == deriv(obj.loss.loss, target[i], x[i])
        end
        ∇input = input_grad(obj)
        @show ∇input
        for i=1:nin
            # TODO a real test
            @test ∇input[i] != 0
        end
    end
end


text:
function KinectV1Cam(;depth_factor=1000.0, baseline=50/525.0)
    cam = RgbdCamParams(
        width=640,
        height=480,
        fx=525.0,
        fy=525.0,
        cx=319.5,
        cy=239.5,
        depth_factor=depth_factor,
        baseline=baseline
    )
    return cam
end

function Aug_ICL_NUIM_CamParam()
    res = KinectV1Cam(depth_factor=1000.0)
    return res
end

"""
TODO: note that TUM have different camera parameters for different sequences
This one is default from ROS bag.
The rule of thumb is to use
"""
function TumCamROS()
    fx = 525.0  # focal length x
    fy = 525.0  # focal length y
    cx = 319.5  # optical center x
    cy = 239.5  # optical center y
    w = 640
    h = 480
    bf = 40.0
    depth_factor = 5000.0
    return RgbdCamParams(
        fx=fx,
        fy=fy,
        cx=cx,
        cy=cy,
        width=w,
        height=h,
        baseline=bf/fx,
        depth_factor=depth_factor)
end


text:
struct Poisson <: Distribution{Int} end

"""
    poisson(lambda::Real)

Sample an `Int` from the Poisson distribution with rate `lambda`.
"""
const poisson = Poisson()

function logpdf(::Poisson, x::Integer, lambda::Real)
    x * log(lambda) - lambda - lgamma(x+1)
end

function logpdf_grad(::Poisson, x::Integer, lambda::Real)
    error("Not implemented")
    (nothing, nothing)
end

function random(::Poisson, lambda::Real)
    rand(Distributions.Poisson(lambda))
end

(::Poisson)(lambda) = random(Poisson(), lambda)
is_discrete(::Poisson) = true

has_output_grad(::Poisson) = false
has_argument_grads(::Poisson) = (false,)

export poisson

max_stars_repo_name: cameronraysmith/Gen

max_issues_repo_name: cameronraysmith/Gen

max_forks_repo_name: cameronraysmith/Gen


text:
# Generate a random orthogonal matrix
function random_orthogonal!{T}(O::Matrix{T}, n::Int)
    F = qrfact!(randn(n, n))
    Q = full(F[:Q])
    R = F[:R]
    for i in 1:size(R, 1)
        d = sign(R[i, i])
        for j in 1:size(Q, 2)
            Q[i, j] *= d
        end
    end

    return copy!(O, Q)
end

random_orthogonal(n::Int) = random_orthogonal!(Matrix{Float64}(n, n), n)

# Generate a random 3D rotation matrix
function random_rotation()
    Q = random_orthogonal(3)
    det(Q) * Q
end

# Generate a random permutation matrix
random_permutation(n::Int) = eye(n)[:, randperm(n)]


text:
# Finish folding the transparent paper according to the instructions. The
# manual says the code is always eight capital letters.
#
# What code do you use to activate the infrared thermal imaging camera system?

# read data
iter = eachline("input.txt")
values = map(l -> parse.(Int, split(l, ',')), Iterators.takewhile(l -> l != "", iter))
folds = map(iter) do l
  (f, i) = split(l, '=')
  (f[end], parse(Int, i))
end

# build sheet - I really hate 1-based indexing
width = maximum(first, values) + 1
height = maximum(last, values) + 1
sheet = falses(width, height)
foreach(((x,y),) -> sheet[x+1,y+1] = true, values)

# fold
for (fold_dir,fold_along) in folds
  if fold_dir == 'x'
    global sheet = sheet[1:fold_along,:] .| sheet[end:-1:(fold_along+2),:]
  else
    global sheet = sheet[:,1:fold_along] .| sheet[:,end:-1:(fold_along+2)]
  end
end

# output - because the code treats sheet as being stored in x,y coordinates and
# julia stores data in column-major form, we can iterate over sheet linearly.
width = size(sheet, 1)
for (i,v) in enumerate(sheet)
  print(v ? '#' : ' ')
  i % width == 0 && print('\n')
end


text:
# Considering every single measurement isn't as useful as you expected: there's
# just too much noise in the data.
#
# Instead, consider sums of a three-measurement sliding window. Again
# considering the above example:
#
# 199  A
# 200  A B
# 208  A B C
# 210    B C D
# 200  E   C D
# 207  E F   D
# 240  E F G
# 269    F G H
# 260      G H
# 263        H
#
# Start by comparing the first and second three-measurement windows. The
# measurements in the first window are marked A (199, 200, 208); their sum is
# 199 + 200 + 208 = 607. The second window is marked B (200, 208, 210); its sum
# is 618. The sum of measurements in the second window is larger than the sum
# of the first, so this first comparison increased.
#
# Your goal now is to count the number of times the sum of measurements in this
# sliding window increases from the previous sum. So, compare A with B, then
# compare B with C, then C with D, and so on. Stop when there aren't enough
# measurements left to create a new three-measurement sum.
#
# In the above example, the sum of each three-measurement window is as follows:
#
# A: 607 (N/A - no previous sum)
# B: 618 (increased)
# C: 618 (no change)
# D: 617 (decreased)
# E: 647 (increased)
# F: 716 (increased)
# G: 769 (increased)
# H: 792 (increased)
# In this example, there are 5 sums that are larger than the previous sum.
#
# Consider sums of a three-measurement sliding window. How many sums are larger
# than the previous sum?

readings = map(l -> parse(Int, l), eachline("input.txt"))
lead = @view readings[2:end]
lead2 = @view readings[3:end]

readings = map(sum, zip(readings, lead, lead2))
lead = @view readings[2:end]

println("Count: ", count(((a, b),) -> a < b, zip(readings, lead)))


text:
import DocStringExtensions
import HTTP
import JSON3

function _add_trailing_slash(url::HTTP.URI)
    _url_string = _get_http_uri_string(url)
    if endswith(_url_string, "/")
        return url
    end
    result = HTTP.URI(string(_url_string, "/"))::HTTP.URI
    return result
end

function _generate_full_url(client::Client,
                            path::AbstractString)
    base_url = get_base_url(client)
    result = _generate_full_url(base_url, path)::HTTP.URI
    return result
end

function _generate_full_url(base_url::BaseURL,
                            path::AbstractString)
    base_url_uri = _get_http_uri(base_url)
    result = _generate_full_url(base_url_uri, path)::HTTP.URI
    return result
end

function _generate_full_url(base_url_uri::HTTP.URI,
                            path::AbstractString)
    base_url_uri_with_trailing_slash = _add_trailing_slash(base_url_uri)
    base_url_uri_string = _get_http_uri_string(base_url_uri_with_trailing_slash)
    full_url_uri_string = string(base_url_uri_string, path)
    full_url_uri = HTTP.URI(full_url_uri_string)::HTTP.URI
    return full_url_uri
end

function _make_http_get_request(client::Client,
                                content_type::V,
                                path::AbstractString) where V <: Val
    http_verb = "GET"
    full_url = _generate_full_url(client, path)
    new_headers = Dict{String, String}()
    content_headers!(new_headers, content_type)
    authentication_headers!(new_headers, client)
    response = HTTP.request(http_verb, full_url, new_headers)
    empty!(new_headers)
    response_body = response.body
    return response_body
end

function get_all_studies(client::Client)
    path = "studies"
    content_type = Val(:json)
    response_body = _make_http_get_request(client, content_type, path)
    response_body_string = String(response_body)::String
    response_body_json = JSON3.read(response_body_string)
    return response_body_json
end

function get_all_series(client::Client;
                        study::AbstractString)
    path = "studies/$(study)/series"
    content_type = Val(:json)
    response_body = _make_http_get_request(client, content_type, path)
    response_body_string = String(response_body)::String
    response_body_json = JSON3.read(response_body_string)
    return response_body_json
end

function get_all_instances(client::Client;
                           study::AbstractString,
                           series::AbstractString)
    path = "studies/$(study)/series/$(series)/instances"
    content_type = Val(:json)
    response_body = _make_http_get_request(client, content_type, path)
    response_body_string = String(response_body)::String
    response_body_json = JSON3.read(response_body_string)
    return response_body_json
end

function get_single_instance_metadata(client::Client;
                                      study::AbstractString,
                                      series::AbstractString,
                                      instance::AbstractString)
    path = "studies/$(study)/series/$(series)/instances/$(instance)/metadata"
    content_type = Val(:json)
    response_body = _make_http_get_request(client, content_type, path)
    response_body_string = String(response_body)::String
    response_body_json = JSON3.read(response_body_string)
    return response_body_json
end

function get_frame_images(::Val{:jpeg},
                          client::Client;
                          study::AbstractString,
                          series::AbstractString,
                          instance::AbstractString,
                          frames::AbstractVector{<:Integer})
    frame_list_string = join(string.(collect(frames)), ",")
    path = "studies/$(study)/series/$(series)/instances/$(instance)/frames/$(frame_list_string)"
    content_type = Val(:jpeg)
    response_body = _make_http_get_request(client, content_type, path)
    images = parse_multipart_image(Val(:jpeg), response_body)
    return images
end


text:
##
using Revise
import Random
using DirectDetections
using Distributions
using DirectImages
using ImageFiltering
using DirectOrbits
using Plots
##

# Generate astrometry points using this template
# and we will try to fit the results
truths = [
    # KeplerianElements(;a = 12,τ = 0.25,ω = 0,Ω = 0,e = 0.1,i = 0.5,μ = 1.,plx = 45.,),
    KeplerianElements(;a = 12,τ = 0.65,ω = 0,Ω = 0,e = 0.1,i = 0.5,μ = 1.,plx = 45.,),
    # KeplerianElements(;a = 18, τ = 0.75,ω = 0,Ω = 0,e = 0.15, i = 0.5,μ = 1.,plx = 45.,),
]
intensities = [
    1800,
    # 3000,
    # 650,
    # 1000,
]

# truth_elements = KeplerianElements(ComponentArray(truth, static))

# times = range(0, period(truth_elements)*4/5, length=9, )
times = range(0, period(first(truths))/2, length=10,)

# Create synthetic images at each time with those points
Random.seed!(1234)
images_contrasts = map(times) do t
    img = 20000fakeimage((201,201),1.2)
    img = centered(img)
    for (inten, truth) in zip(intensities, truths)
        ra, dec = kep2cart(truth, t)
        x = -ra
        y = dec
        img[round(Int,x/10), round(Int,y/10)] += inten #2500 #+ 300randn()
    end

    img = imfilter(img, Kernel.gaussian(5), "replicate")
    contrast = contrast_interp(img)

    img, contrast
end
images = [img for (img,contrast) in images_contrasts]
contrasts = [contrast for (img,contrast) in images_contrasts]
nothing

##
# imshow2(reduce(hcat, images), clims=(-35,35))

##

planet_priors = Priors(
    a = Uniform(7, 16),
    e = TruncatedNormal(0.1, 0.1, 0.0, 0.4),
    τ = Uniform(0,1),
    # τ = Uniform(0.4,0.8),
    ω = Uniform(-π,π),
    # ω = Uniform(-π/4,π/4),
    i = Uniform(-π,π),
    # i = Uniform(-π/4,π/4),
    # i = TruncatedNormal(0, 0.2, -π,π),
    Ω = Uniform(-π,π),
    # Ω = Uniform(-π/4,π/4),

    # mass = Uniform(0mjup2msol,70mjup2msol),

    J = Uniform(0,20)
)
planet = reparameterize(DirectDetections.Planet(planet_priors))
nothing
##

system_images = DirectDetections.Images(
    (band=:J, image=images[1], platescale=10.0, epoch=times[1]),#, contrast=contrasts[1]),
    (band=:J, image=images[2], platescale=10.0, epoch=times[2]),#, contrast=contrasts[2]),
    (band=:J, image=images[3], platescale=10.0, epoch=times[3]),#, contrast=contrasts[3]),
    (band=:J, image=images[4], platescale=10.0, epoch=times[4]),#, contrast=contrasts[4]),
)

system_priors = Priors(
    μ = TruncatedNormal(1.0, 0.01, 0, 10),
    plx = TruncatedNormal(45., 0.0001, 0, 100),
)

# system = System(system_priors, system_pma, system_images, planet_b)
system = System(system_priors, system_images, planet,)
nothing

##
@time chains_img = DirectDetections.mcmc(
    system;
    # numwalkers=800,
    # burnin=62_000,
    # numsamples_perwalker=100_000,
    # thinning=250,

    numwalkers=50,
    burnin=1000,
    numsamples_perwalker=5_000,
    thinning=1,

    squash = false
);
nothing

##
@time chainsh, stats = DirectDetections.hmc(
    system;
    burnin=10_000,
    numwalkers=10,
    numsamples_perwalker=15_000,
);
nothing

##
[count(getproperty.(s, :numerical_error)) for s in stats]
##
function orbsonimg(args...; kwargs...)
    plot()
    orbsonimg!(args...;kwargs...)
end
function orbsonimg!(planet, image, N=100; kwargs...)
    sampled = sampleorbits(planet, N);
    i = DirectImage(image)
    i.PLATESCALE = 10.
    imshow!(i;skyconvention=true,kwargs...)
    # plot!(sampled,color=:white,alpha=5/N,label="")
    plot!(sampled,color=:white,label="",)
    xlims!(i.PLATESCALE.*extrema(axes(i,1)))
    ylims!(i.PLATESCALE.*extrema(axes(i,2)))

end

##
imshow(images[3],clims=(-2,10))
# orbsonimg(chains.planets[1],images[2],clims=(-2,10))
orbsonimg(chains.planets[1],mean(images), 500,legend=nothing, clims=(-0.2,1),color=:magma,dpi=200)
plot!(truths, color=:black, lw=2, ls=:dash, label="")
# orbsonimg(chainsh[2].planets[1],images)
# writefits("tmp.fits", images[2])
##
plot()
orbsonimg!(chainsh[1].planets[1],mean(images),1,legend=nothing, clims=(-0.2,1),color=:magma,dpi=200,)
for c in chainsh
    plot!(sampleorbits(c.planets[1],125),color=:white,label="",alpha=0.1)
end
    plot!(truths, color=:black, lw=2, ls=:dash, label="",colorbar=nothing)

##

function snr(phot)
    low,med,up = quantile(reshape(phot,:),(0.33, 0.5, 0.67))
    st = std(reshape(phot,:))
    return med/(up-low)
end

##
for c in chainsh
    snr(c.planets[1].phot.Keck_L′) |> println
end
##
snr(chains.planets[1].phot.Keck_L′[:])
##
histogram(chains.planets[1].phot.Keck_L′[:])
histogram(chainsh[1].planets[1].phot.Keck_L′[:])
##
DirectDetections.plotposterior(chains_img, 1, :a, 500)
plot!(truths, color=:black, lw=2, ls=:dash, label="", alpha=0.5)
# savefig("images/readme-orbits.png")

##
DirectDetections.plotposterior(chains_img, chains_img.planets[1], :i, 500, colorbartitle="inclination (rad)",dpi=200)
# savefig("images/readme-post-i.png")

##
DirectDetections.plotposterior(chains_img, chains_img.planets[1], (:phot, :Keck_L′), 500, colorbartitle="flux", cmap=:plasma, rev=false, clims=(0,12),dpi=200)
plot!(truths, color=:black, lw=2, ls=:dash, label="")
# savefig("images/readme-post-f.png")

##
DirectDetections.plotposterior(chains.planets[1], :e, 500, colorbartitle="eccentricity",dpi=200)
plot!(truths, color=:black, lw=2, ls=:dash)
# savefig("images/readme-post-e.png")

##
DirectDetections.plotposterior(chains.planets[1],  lw=3, :e, 500, colorbartitle="eccentricity",dpi=200, cmap=:plasma, rev=false,)
plot!(truths, color=:black, lw=2, ls=:dash)
# savefig("images/readme-post-e2.png")

##

ra, dec = projectpositions(chains.planets[1], mean(times))
histogram2d(ra,dec,aspectratio=1, color=:plasma, background_inside=:black, framestyle=:box,xflip=true,dpi=200, colorbartitle=".\n\nposterior density")
xlims!(1.25.*(-1,+1).*max(abs.(extrema(ra))...))
ylims!(1.25.*(-1,+1).*max(abs.(extrema(dec))...))
xlabel!("Δ right ascension (as)")
ylabel!("Δ declination (as)")

##
mask = (-200 .< ra .< 250) .& (300 .< dec .< 800)
snr(chains.planets[1].phot.Keck_L′)
snr(chains.planets[1].phot.Keck_L′[mask])
histogram2d(ra[mask],dec[mask],aspectratio=1, color=:plasma, background_inside=:black, framestyle=:box, cscale=:log10kep2cart,dpi=200)
histogram2d(ra[mask],dec[mask],aspectratio=1, color=:plasma, background_inside=:black, framestyle=:box, cscale=:log10kep2cart,dpi=200)
# xlims!(1.25.*(-1,+1).*max(abs.(extrema(ra))...))
# ylims!(1.25.*(-1,+1).*max(abs.(extrema(dec))...))
##
ras,decs = projectpositions(chains.planets[1],mean(times))
# histogram2d(ras[5end÷6:end], decs[5end÷6:end])
histogram2d(ras, decs,aspectratio=1,xflip=true,color=:plasma)

##
mask = 7 .< chains.planets[1].a .< 15
snr(chains.planets[1].phot.Keck_L′[mask])
mask = chains.planets[1].a .< 7
snr(chains.planets[1].phot.Keck_L′[mask])

##
table = (;
    L=chains.planets[1].phot.Keck_L′,
    chains.planets[1].a,
    chains.planets[1].i,
    chains.planets[1].e,
    tau=chains.planets[1].τ,
)
corner(
    table,
    split(raw"\mathrm{flux} a i e \tau"),
    plotscatter=false,
    hist2d_kwargs=(;ylims=(0,NaN),xlims=(0,NaN)),
    scatter_kwargs=(;markersize=0.5),
    dpi=200,
    # lim_factor=2,
    # hist2d_kwargs=(;nbins=8),
)

##
ra, dec = projectpositions(chains.planets[1], mean(times))
pa = atan.(dec,ra)
# mask = (pa .< 0.5) .&
#     (600 .< ra .< 800) .&
#     (-200 .< dec .< -100)
mask = vec((chains.planets[1].phot.Keck_L′ .< 6))
table = (;
    ra=ra[mask],
    dec=dec[mask],
    # pa=rad2deg.(pa[mask]),
    a=(chains.planets[1].a)[mask],
    L=(chains.planets[1].phot.Keck_L′)[mask],
    i=rad2deg.((chains.planets[1].i))[mask],
    e=(chains.planets[1].e)[mask],
)
corner(
    table,
    # split(raw"a \mathrm{flux} i e \tau"),
    plotscatter=false,
    # hist2d_kwargs=(;ylims=(0,NaN),xlims=(0,NaN)),
    # scatter_kwargs=(;markersize=0.5),
    # dpi=200,
    # lim_factor=2,
    hist_kwargs=(;nbins=50),
    hist2d_kwargs=(;nbins=50),
)

## Chains
table = (;
    a=(chains.planets[1].a),
    L=(chains.planets[1].phot.Keck_L′),
    i=rad2deg.((chains.planets[1].i)),
    e=(chains.planets[1].e),
)
corner(
    table,
    # split(raw"a \mathrm{flux} i e \tau"),
    plotscatter=false,
    # hist2d_kwargs=(;ylims=(0,NaN),xlims=(0,NaN)),
    # scatter_kwargs=(;markersize=0.5),
    # dpi=200,
    # lim_factor=2,
    hist_kwargs=(;nbins=50),
    hist2d_kwargs=(;nbins=50),
)

## Chainsh
table = (;
    a=reduce(vcat, [c.planets[1].a for c in chainsh]),
    L=reduce(vcat, [c.planets[1].phot.Keck_L′ for c in chainsh]),
    i=rad2deg.(reduce(vcat, [c.planets[1].i for c in chainsh])),
    e=reduce(vcat, [c.planets[1].e for c in chainsh]),
)
corner(
    table,
    # split(raw"a \mathrm{flux} i e \tau"),
    plotscatter=false,
    # hist2d_kwargs=(;ylims=(0,NaN),xlims=(0,NaN)),
    # scatter_kwargs=(;markersize=0.5),
    # dpi=200,
    # lim_factor=2,
    hist_kwargs=(;nbins=25),
    hist2d_kwargs=(;nbins=25),
)


text:
using ComponentArrays
using ControlSystems
using DifferentialEquations
using UnPack
using Plots

# ## Helper functions
# First, we need a way to apply inputs to the system through keyword arguments. These
# will help us pass in inputs as either values or functions of (x,p,t).
maybe_apply(f::Function, x, p, t) = f(x, p, t)
maybe_apply(f, x, p, t) = f

function apply_inputs(func; kwargs...)
    simfun(dx, x, p, t) = func(dx, x, p, t; map(f->maybe_apply(f, x, p, t), (;kwargs...))...)
    simfun(x, p, t) = func(x, p, t; map(f->maybe_apply(f, x, p, t), (;kwargs...))...)
    return simfun
end

# Next, we need a way to create derivative functions from transfer functions. In ControlSystems.jl
# there is a function called `simulator` that does this, but the inputs must be applied from
# the start, so we couldn't use it as a component function. Our version allows inputs to be
# passed through the keyword arguments and, as an added convenience, is in observer canonical
# form so our first element of `x` is also the output `y` (note that while this is true for
# our problem, it isn't always going to be the case).
SISO_simulator(P::TransferFunction) = SISO_simulator(ss(P))
function SISO_simulator(P::AbstractStateSpace)
    @unpack A, B, C, D = P

    if size(D)!=(1,1)
        error("This is not a SISO system")
    end

    # Put into observer canonical form so the first element is also the y value
    BB = reverse(vec(C))
    CC = reverse(vec(B))'
    DD = D[1,1]

    return function sim!(dx, x, p, t; u=0.0)
        dx .= A*x + BB*u
        return CC*x + DD*u
    end
end

## Model setup
# We'll make a Laplace variable `s`
s = tf("s")

# Here is our reference model to track
am = 3
bm = 3
ref_model = bm / (s + am)
ref_sim! = SISO_simulator(ref_model)

# and our plant model. The nominal plant structure is what is known to our adaptation law.
ap = 1
bp = 2
nominal_plant = bp / (s + ap)
nominal_sim! = SISO_simulator(nominal_plant)

# To test robustness to uncertainty, we'll also include unmodeled dynamics with an entirely
# different structure than our nominal plant model.
unmodeled_dynamics = 229/(s^2 + 30s + 229)
truth_plant = nominal_plant * unmodeled_dynamics
truth_sim! = SISO_simulator(truth_plant)

# We'll make a first-order sensor as well so we can add noise to our measurement
τ = 0.005
sensor_plant = 1 / (τ*s + 1)
sensor_sim! = SISO_simulator(sensor_plant)

## Derivative functions
# Our control law assumes perfect knowledge of the parameters that are attached to the
# regressors (which are the reference input and the model output)
control(θ, w) = θ'w

# We'll use a simple gradient descent adaptation law
function adapt!(Dθ, θ, γ, t; e, w)
    Dθ .= -γ*e*w
    return nothing
end

# Our feedback loop takes in the reference model output `ym` and the input signal `r`,
# calculates the control signal `u`, feeds that into the plant model, calculates the reference
# tracking error `e`, and finally updates feeds the reference tracking error and it's corresponding
# regressor vector to the adaptation law.
function feedback_sys!(D, vars, p, t; ym, r, n)
    @unpack parameter_estimates, plant_model, sensor = vars
    γ = p.gamma
    regressor = [r, plant_model[1]]

    u = control(parameter_estimates, regressor)
    yp = p.plant_fun(D.plant_model, plant_model, (), t; u=u)
    ŷ = sensor_sim!(D.sensor, sensor, (), t; u=yp[1]) + n
    e = ŷ .- ym
    regressor[2] = ŷ
    adapt!(D.parameter_estimates, parameter_estimates, γ, t; e=e, w=regressor)
    return yp
end
# Now the full system takes in an input signal `r`, feeds it through the reference model,
# and feeds the output of the reference model `ym` and the input signal to `feedback_sys`.
function system!(D, vars, p, t; r=0.0, n=0.0)
    @unpack reference_model, feedback_loop = vars

    ym = ref_sim!(D.reference_model, reference_model, (), t; u=r)
    yp = feedback_sys!(D.feedback_loop, feedback_loop, p, t; ym=ym, r=r, n=n)
    return yp
end

## Simulation inputs
# Simulation time span
tspan = (0.0, 30.0)

# Input signal
input_signal = (x,p,t) -> sin(3t)

# Initial conditions
ref_ic = zeros(1)
nominal_ic = zeros(1)
truth_ic = zeros(3)
sensor_ic = zeros(1)
θ_est_ic = ComponentArray(θr=0.0, θy=0.0)

## Set up and run Simulation
function simulate(plant_fun, plant_ic;
                tspan=tspan,
                input_signal=input_signal,
                adapt_gain=1.5,
                noise_param=nothing,
                deterministic_noise=0.0)

    noise(D, vars, p, t) = (D.feedback_loop.sensor[1] = noise_param)

    # Truth control parameters
    θ_truth = (r=bm/bp, y=(ap-am)/bp)

    # Initial conditions
    ic = ComponentArray(
        reference_model = ref_ic,
        feedback_loop = (
            parameter_estimates = θ_est_ic,
            sensor = sensor_ic,
            plant_model = plant_ic,
        ),
    )

    # Model parameters
    p = (
        gamma = adapt_gain,
        plant_fun = plant_fun,
    )

    sim_fun = apply_inputs(system!; r=input_signal, n=deterministic_noise)

    # We can also choose whether we want to include noise in our model by switching between an ODE
    # and an SDE problem.
    if noise_param === nothing
        prob = ODEProblem(sim_fun, ic, tspan, p, max_iters=2000)
    else
        prob = SDEProblem(sim_fun, noise, ic, tspan, p, max_iters=2000) # With noise
    end

    ## Solve!
    sol = solve(prob)

    ## Plot
    # Reference model tracking
    top = plot(
        sol,
        vars=Symbol.(["reference_model[1]", "feedback_loop.sensor[1]"]),
        # legend=:right,
        title="Reference Model Tracking",
    )

    # Parameter estimate tracking
    bottom = plot(sol, vars=Symbol.(["feedback_loop.parameter_estimates.θr", "feedback_loop.parameter_estimates.θy"]))
    plot!(
        bottom,
        [tspan...], [θ_truth.r θ_truth.y; θ_truth.r θ_truth.y],
        labels=["θr truth" "θy truth"],
        legend=:right,
        title="Parameter Estimate Tracking",
    )

    # Combine both plots
    plot(top, bottom, layout=(2,1), size=(800, 800))
end

simulate(truth_sim!, truth_ic; input_signal=2.0, deterministic_noise=(x,p,t)->0.5sin(16.1t), noise_param=nothing)

max_forks_repo_name: jonniedie/ComponentArrays


text:
using LuckyImaging
using Documenter

DocMeta.setdocmeta!(LuckyImaging, :DocTestSetup, :(using LuckyImaging); recursive=true)

makedocs(;
    modules=[LuckyImaging],
    authors="Miles Lucas <mdlucas@hawaii.edu> and contributors",
    repo="https://github.com/JuliaHCI/LuckyImaging.jl/blob/{commit}{path}#{line}",
    sitename="LuckyImaging.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "Examples" => "examples.md",
        "API/Reference" => "api.md"
    ],
)

deploydocs(;
    devbranch="main",
    repo="github.com/JuliaHCI/LuckyImaging.jl",
)