% 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
omegad = -2;

t0 = 1;
q0 = [0 pi/2 -pi/4 0 -pi/2 0]';
dq0 = zeros(6,1);
d2q0 = zeros(6,1);

t1 = 2.25;
q1 = [0 pi/4 -pi/2 0 -pi/2 0]';
dq1 = [0 omegad omegad 0 0 0]';

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


% 8 conditions -- use 7th order poly
A =@(t) [1 t t^2 t^3 t^4 t^5  t^6 t^7];
dA =@(t) [0 1 2*t 3*t^2 4*t^3 5*t^4 6*t^5 7*t^6];
d2A =@(t) [0 0 2 6*t 12*t^2 20*t^3 30*t^4 42*t^5];

b = [q0'; dq0'; d2q0'; q1'; dq1'; qf'; dqf'; d2qf'];
b = b(:);

Ai =@(t0,t1,tf) [A(t0); dA(t0); d2A(t0);
                 A(t1); dA(t1);
                 A(tf); dA(tf); d2A(tf)];
A = blkdiag(Ai(t0,t1,tf),Ai(t0,t1,tf),Ai(t0,t1,tf), ...
            Ai(t0,t1,tf),Ai(t0,t1,tf),Ai(t0,t1,tf));
a = A\b;

% make joint posistion and velocity functions
q_poly =@(t) [polyval(flipud(a(1:8)),t);
              polyval(flipud(a((1:8)+8)),t);
              polyval(flipud(a((1:8)+2*8)),t);
              polyval(flipud(a((1:8)+3*8)),t);
              polyval(flipud(a((1:8)+4*8)),t);
              polyval(flipud(a((1:8)+5*8)),t)];
fqd =@(t) (t>=t0 & t<=tf)*q_poly(t) + (t<t0)*q0 + (t>tf)*qf;
qd = arrayfun(@(t) fqd(t)',t,'UniformOutput',0);
qd = cell2mat(qd);

dq_poly =@(t) [polyval(polyder(flipud(a(1:8))),t);
              polyval(polyder(flipud(a((1:8)+8))),t);
              polyval(polyder(flipud(a((1:8)+2*8))),t);
              polyval(polyder(flipud(a((1:8)+3*8))),t);
              polyval(polyder(flipud(a((1:8)+4*8))),t);
              polyval(polyder(flipud(a((1:8)+5*8))),t)];
fdqd =@(t) (t>=t0 & t<=tf)*dq_poly(t) + (t<t0)*dq0 + (t>tf)*dqf;
dqd = arrayfun(@(t) fdqd(t)',t,'UniformOutput',0);
dqd = cell2mat(dqd);

Fqd =@(t) [fqd(t); fdqd(t)];

% Find maximum velocity
[~,idx] = max(abs(dqd(:,2)));
tfly = t(idx);


figure;
  subplot(211);
  plot(t,qd); hold all;
  plot([t0,t1,tf],b([1,4,6] + 8),'o');
  plot([tfly,tfly],ylim,'--k');
  ylabel('q_2');
  subplot(212);
  plot(t,qd(:,3)); hold all;
  plot([t0,t1,tf],b([1,4,6] + 2*8),'o');
  plot([tfly,tfly],ylim,'--k');
  ylabel('q_3');

figure;
  subplot(211);
  plot(t,dqd(:,2)); hold all;
  plot([t0,t1,tf],b([2,5,7] + 8),'o');
  plot([tfly,tfly],ylim,'--k');
  ylabel('dq_2');
  subplot(212);
  plot(t,dqd(:,3)); hold all;
  plot([t0,t1,tf],b([2,5,7] + 2*8),'o');
  plot([tfly,tfly],ylim,'--k');
  ylabel('dq_3');

%figure;
%  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 = 500;
kd = 100;

% Ball
rad = 0.05;
m = 0.1;
[xb,yb,zb] = sphere;
xb = rad*xb;
yb = rad*yb;
zb = rad*zb;

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

% Dynamics with ball attached
fext =@(t,q) externalForce(r,'tool0',[0 0 0 0 0 -m*9.81]',q(1:n));
f1 =@(t,q) [q((1:n)+n); forwardDynamics(r, ...
                           q(1:n), ...
                           q((1:n)+n), ...
                           u(t,q,Fqd(t)), ...
                           fext(t,q))];

% Dynamics without ball attached
fball =@(t,x) [x(4:end); ...
              0; 0; -m*9.81];
f2 =@(t,q) [q((1:n)+n); forwardDynamics(r, ...
                          q(1:n), ...
                          q((1:n)+n), ...
                          u(t,q,Fqd(t)), ...
                          []); ...
                          fball(t,q((2*n+1):end))];

tic
ifly = find(t == tfly);

% First sim
[t1,q1] = ode23tb(@(t,q) f1(t,q),t(1:ifly),[q0;q0]);

% Store end-effector position (ball position) during first part
rb = zeros(numel(1:ifly),3);
for i=1:numel(t(1:ifly))
  ee = getTransform(r,q1(i,1:n)','tool0');
  rb(i,:) = ee(1:3,4);
end

% Get initial conditions for next sim
v0 = geometricJacobian(r,q1(end,1:n)','tool0')*q1(end,(1:n)+n)';
x0 = [rb(end,:)'; v0(4:end)];

% Second Sim
[t2,q2] = ode23tb(@(t,q) f2(t,q),t(ifly:end),[q1(end,:)';x0]);
q = [q1; q2(2:end,1:2*n)];
rb = [rb; q2(2:end,13:15)];

toc

% Kinematics
figure;
  ax = show(r,q0, ...
           'Visuals','on', ...
           'PreservePlot',0, ...
           'Fastupdate',1); hold all;
  ball = surf(xb + rb(1,1),yb + rb(1,2),zb + rb(1,3));
  ball.EdgeColor = 'k';
  ball.FaceColor = 'k';
  drawFloor();
for i=1:10:numel(t)
  show(r,q(i,1:n)', ...
        'Parent',ax, ...
        'PreservePlot',0, ...
        'Fastupdate',1);
  ball.XData = xb + rb(i,1);
  ball.YData = yb + rb(i,2);
  ball.ZData = zb + rb(i,3);
  pause(0.01);
end

figure;
  plot(t,qd,'--'); hold all;
  plot(t,q(:,1:n));
  xlabel('time (s)')
  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([-1.75,0.75]);
  ylim([-0.75,0.75]);
  zlim([0,1]);
end


