% 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 and contraints
t0 = 1;
q0 = zeros(1,n)';
dq0 = zeros(1,n)';
d2q0 = zeros(1,n)';


tf = 3;
qf = [0 pi/2 -pi/2 0 pi/4 0]';
dqf = zeros(1,n)';
d2qf = zeros(1,n)';


% -----------------------------------------------------------------------------
% quintic polynomial in joint space
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];

Ai =@(t0,tf) [A(t0); dA(t0); d2A(t0);
              A(tf); dA(tf); d2A(tf)];

% make system of linear equations
b = [q0'; dq0'; d2q0'; qf'; dqf'; d2qf'];
b = b(:);
A = blkdiag(Ai(t0,tf),Ai(t0,tf),Ai(t0,tf),Ai(t0,tf),Ai(t0,tf),Ai(t0,tf));
a = A\b;

% make joint angle function
q_poly =@(t) [polyval(flipud(a(1:n)),t);
              polyval(flipud(a((1:n)+n)),t);
              polyval(flipud(a((1:n)+2*n)),t);
              polyval(flipud(a((1:n)+3*n)),t);
              polyval(flipud(a((1:n)+4*n)),t);
              polyval(flipud(a((1:n)+5*n)),t)];
% poly only applies between t0 tf
q_traj =@(t) (t>=t0 & t<=tf)*q_poly(t) + (t<t0)*q0 + (t>tf)*qf;

dq_poly =@(t) [polyval(polyder(flipud(a(1:n))),t);
              polyval(polyder(flipud(a((1:n)+n))),t);
              polyval(polyder(flipud(a((1:n)+2*n))),t);
              polyval(polyder(flipud(a((1:n)+3*n))),t);
              polyval(polyder(flipud(a((1:n)+4*n))),t);
              polyval(polyder(flipud(a((1:n)+5*n))),t)];
dq_traj =@(t) (t>=t0 & t<=tf)*dq_poly(t) + (t<t0)*dq0 + (t>tf)*dqf;


%% Simulate trajectory only
figure('Position',[0 0 560 420]);
  ax = show(r,q0, ...
           'Visuals','on', ...
           'PreservePlot',0, ...
           'Fastupdate',1); hold all;
  drawFloor();
  title('Trajectory Only');

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

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

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

dynamics =@(t,q) [q((1:n)+n); forwardDynamics(r, ...
                                  q(1:n), ...
                                  q((1:n)+n), ...
                                  tau(t,q,b,q_traj(t),dq_traj(t),kp,kd), ...
                                  [])];

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

%% Simulate controller
figure('Position',[0 0 560 420]);
  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

qd = arrayfun(@(t) q_traj(t)',t,'UniformOutput',0);
qd = cell2mat(qd);

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


% -----------------------------------------------------------------------------
function drawFloor()
  ax = gca;
  ax.CameraViewAngle = 4;
  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
