using LinearAlgebra
using RigidBodyDynamics
using StaticArrays

# gravitational acceleration
g = -9.81

# link lengths
l_0 = 1.10
l_1 = 0.5
l_2 = 1.20
l_3 = 0.75

# link masses
m_1 = 0.5
m_2 = 1.0
m_3 = 0.75

# link center of mass offsets from the preceding joint axes
c_1 = 0.25
c_2 = 0.60
c_3 = 0.375

# moments of inertia about the center of mass of each link
I_1 = 0.333
I_2 = 0.537
I_3 = 0.4

# scalar type
T = Float64;

# Rotation axis: negative y-axis
axis = SVector(zero(T), -one(T), zero(T));

world = RigidBody{T}("world")
fourbar = Mechanism(world; gravity = SVector(0., 0., g))

# link1 and joint1
joint1 = Joint("joint1", Revolute(axis))
inertia1 = SpatialInertia(CartesianFrame3D("inertia1_centroidal"), I_1*axis*axis', zero(SVector{3, T}), m_1)
link1 = RigidBody(inertia1)
before_joint1_to_world = one(Transform3D, frame_before(joint1), default_frame(world))
c1_to_joint = Transform3D(inertia1.frame, frame_after(joint1), SVector(c_1, 0, 0))
attach!(fourbar, world, link1, joint1, joint_pose = before_joint1_to_world, successor_pose = c1_to_joint)

# link2 and joint2
joint2 = Joint("joint2", Revolute(axis))
inertia2 = SpatialInertia(CartesianFrame3D("inertia2_centroidal"), I_2*axis*axis', zero(SVector{3, T}), m_2)
link2 = RigidBody(inertia2)
before_joint2_to_after_joint1 = Transform3D(frame_before(joint2), frame_after(joint1), SVector(l_1, 0., 0.))
c2_to_joint = Transform3D(inertia2.frame, frame_after(joint2), SVector(c_2, 0, 0))
attach!(fourbar, link1, link2, joint2, joint_pose = before_joint2_to_after_joint1, successor_pose = c2_to_joint)

# link3 and joint3
joint3 = Joint("joint3", Revolute(axis))
inertia3 = SpatialInertia(CartesianFrame3D("inertia3_centroidal"), I_3*axis*axis', zero(SVector{3, T}), m_3)
link3 = RigidBody(inertia3)
before_joint3_to_world = Transform3D(frame_before(joint3), default_frame(world), SVector(l_0, 0., 0.))
c3_to_joint = Transform3D(inertia3.frame, frame_after(joint3), SVector(c_3, 0, 0))
attach!(fourbar, world, link3, joint3, joint_pose = before_joint3_to_world, successor_pose = c3_to_joint)

# joint between link2 and link3
joint4 = Joint("joint4", Revolute(axis))
before_joint4_to_joint2 = Transform3D(frame_before(joint4), frame_after(joint2), SVector(l_2, 0., 0.))
joint3_to_after_joint4 = Transform3D(frame_after(joint3), frame_after(joint4), SVector(-l_3, 0., 0.))
attach!(fourbar, link2, link3, joint4, joint_pose = before_joint4_to_joint2, successor_pose = joint3_to_after_joint4)

state = MechanismState(fourbar)
result = DynamicsResult(fourbar);

set_configuration!(state, joint1, 1.6707963267948966) # θ
set_configuration!(state, joint2, -1.4591054166649482) # γ
set_configuration!(state, joint3, 1.5397303602625536) # ϕ

set_velocity!(state, joint1, 0.5)
set_velocity!(state, joint2, -0.47295)
set_velocity!(state, joint3, 0.341)

# Invalidate the cache variables
setdirty!(state)

ts, qs, vs = simulate(state, 3., Δt = 1e-2);

#NBSKIP
using RigidBodyTreeInspector

#NBSKIP
DrakeVisualizer.any_open_windows() || (DrakeVisualizer.new_window(); sleep(1));

#NBSKIP
vis = DrakeVisualizer.Visualizer()[:robot1];
RigidBodyTreeInspector.setgeometry!(vis, fourbar, visual_elements(fourbar, Skeleton(inertias = false)));

#NBSKIP
RigidBodyTreeInspector.animate(vis, fourbar, ts, qs; realtimerate = 1.);