% Robot setup
r = loadrobot('abbIrb120','DataFormat','column');
r.Gravity = [0 0 -9.81];
n = 6;

% time vector
dt = 1/1000;
T = 5;
t = 0:dt:(T-dt);
t = t(:);

% Trajectory end points
t0 = 1;
q0 = zeros(1,n)';
T0 = getTransform(r,q0,'tool0');

tf = 3;
qf = [0 pi/2 -pi/2 0 pi/4 0]';
Tf = getTransform(r,qf,'tool0');

% Cartesian planning with quintic time scaling
p0 = T0(1:3,4);
pf = Tf(1:3,4);
p =@(s) p0 + s*(pf - p0);

R0 = tform2rotm(T0);
Rf = tform2rotm(Tf);
R =@(s) R0*expm(logm(R0'*Rf)*s);

A =@(t) [1 t t^2 t^3 t^4 t^5];
dA =@(t) [0 1 2*t 3*t^2 4*t^3 5*t^4];
d2A =@(t) [0 0 2 6*t 12*t^2 20*t^3];

b = [0; 0; 0; 1; 0; 0];
A = [A(t0); dA(t0); d2A(t0); A(tf); dA(tf); d2A(tf)];
a = inv(A)*b;
s =@(t) (t>=1 & t<=3)*polyval(flipud(a),t) + (t<1)*0 + (t>3)*1;
Td =@(t) trvec2tform(p(s(t))')*rotm2tform(R(s(t)));

aik = analyticalInverseKinematics(r);
generateIKFunction(aik,'robotIK');

% IK - Generate joint trajectories from end-effector trajectories
qd = zeros(numel(t),n);
for i=1:(numel(t)-1)
  qd(i+1,:) = myIK(Td(t(i)),qd(i,:)')';
end
qd =@(tq) interp1(t,qd,tq)';

%% Simulate trajectory tracking with PD
b  = 0.1;
kp = 1000;
kd = 100;

% joint torques -- damping + PD controller
u =@(q,qd) -b*q((1:n)+n) ...
              + (qd-q(1:n))*kp ...
              + (-q((1:n)+n))*kd;

f =@(t,q) [q((1:n)+n); forwardDynamics(r, ...
                                  q(1:n), ...
                                  q((1:n)+n), ...
                                  u(q,qd(t)), ...
                                  [])];

tic
[t,q] = ode23tb(@(t,q) f(t,q),t,[q0;q0]);
toc

% Animation
figure;
  ax = show(r,q0, ...
           'Visuals','on', ...
           'PreservePlot',0, ...
           'Fastupdate',1); hold all;
  drawFloor();
  title('PD-Controller Tracking');

for i=1:10:numel(t)
  show(r,q(i,1:n)', ...
        'Parent',ax, ...
        'PreservePlot',0, ...
        'Fastupdate',1);
  pause(0.01);
end


figure;
  plot(t,qd(t),'--'); hold all;
  plot(t,q(:,1:n));
  xlabel('time (s)')
  ylabel('q & q_d')

% -----------------------------------------------------------------------------
function drawFloor()
  ax = gca;
  ax.CameraViewAngle = 5;
  p = patch([1 -1 -1 1].*0.5, [1 1 -1 -1]*0.5, [0 0 0 0]);
  p.FaceColor = [0.8, 0.8, 0.8];
  axis off;
  xlim([-0.75,0.75]);
  ylim([-0.75,0.75]);
  zlim([0,1]);
end

function q = myIK(T,q)
  q_possible = robotIK(T,1,1,q);
  [~,idx] = min(vecnorm((q_possible - q')'));
  q = q_possible(idx,:)';
end

