r = loadrobot('abbIrb120','DataFormat','column');
n = 6; 
q0 = zeros(6,1);

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

%% ----------------------------------------------------------------------------
% Generalized IK with cup

% Add a cup
cupHeight = 0.2;
cupRadius = 0.05;
xd = 0.4;
yd = 0;
cupPosition = [xd, yd, 0];
cup = rigidBody('cupFrame');

cupTargetPos = [-yd, xd, 0];
cupTarget = rigidBody('cupTargetFrame');

graspOffset = [0 -0.05 cupHeight*0.6];

% Scale visual
addVisual(cup,"Mesh",'cup.stl',[[0.0015*eye(3), zeros(3,1)]; 0 0 0 1]);

setFixedTransform(cup.Joint,trvec2tform(cupPosition))
addBody(r,cup,r.BaseName);

setFixedTransform(cupTarget.Joint,trvec2tform(cupTargetPos))
addBody(r,cupTarget,r.BaseName);

% Generalized IK object
gik = generalizedInverseKinematics('RigidBodyTree',r);
gik.ConstraintInputs = {'cartesian','position','aiming'};

% Add Floor Constraint
heightAboveFloor = constraintCartesianBounds('tool0');
heightAboveFloor.Bounds = [-inf, inf; ...
                           -inf, inf; ...
                           0.005, inf];

% Add Position Constraint
distanceFromCup = constraintPositionTarget('tool0');
distanceFromCup.ReferenceBody = 'cupFrame';
distanceFromCup.TargetPosition = graspOffset;

% Aim at the cup
alignWithCup = constraintAiming('tool0');
alignWithCup.TargetPoint = [0 0 100];

% Get IK
[q1,solnInfo] = gik(q0,heightAboveFloor, ...
                      distanceFromCup, ...
                      alignWithCup);

% Add Position Constraint
distanceFromCup = constraintPositionTarget('tool0');
distanceFromCup.ReferenceBody = 'cupTargetFrame';
distanceFromCup.TargetPosition = [0 -0.05 cupHeight*0.6];

% Aim at the cup
alignWithCup = constraintAiming('tool0');
alignWithCup.TargetPoint = [0 0 100];

% Get IK
[q2,solnInfo] = gik(q0,heightAboveFloor, ...
                      distanceFromCup, ...
                      alignWithCup);

% Plan trajectory
t0 = 1;
dq0 = zeros(n,1);
d2q0 = zeros(n,1);

t1 = 2;
dq1 = zeros(n,1);
d2q1 = zeros(n,1);

t2 = 3;
dq2 = zeros(n,1);
d2q2 = zeros(n,1);

t3 = 4;

[q01, dq01] = quinticPoly(t0,q0,dq0,d2q0,t1,q1,dq1,d2q1);
[q12, dq12] = quinticPoly(t1,q1,dq1,d2q1,t2,q2,dq2,d2q2);
[q20, dq20] = quinticPoly(t2,q2,dq2,d2q2,t3,q0,dq0,d2q0);

qd =@(t) (t<=t1)*q01(t) + (t>t1 & t<=t2)*q12(t) + (t>t2)*q20(t);

%v = VideoWriter('pick-n-place-IK-only.mp4','MPEG-4');
%open(v);
figure;
  ax = show(r,q0, ...
    'Visuals','on', ...
    'PreservePlot',0, ...
    'Fastupdate',1); hold all;
  scatter3(cupTargetPos(1), ...
           cupTargetPos(2),cupTargetPos(3)+0.1,100,'r','filled');
  drawFloor();

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

  % Simuate the 'grasp' by changing the position of the cup during the 'place'
  % motion
  Ti = getTransform(r,qd(t(i)),'tool0');
  ee = Ti(1:3,4);
  cupPos = (t(i)<=t1)*cupPosition ...
           + (t(i)>t1 & t(i)<=t2)*(ee' - graspOffset) ...
           + (t(i)>t2)*cupTargetPos;
  setFixedTransform(r.Bodies{9}.Joint,trvec2tform(cupPos));

%  frame = getframe(gca);
%  writeVideo(v,frame);
  pause(0.01);
end
%close(v);

% -----------------------------------------------------------------------------
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 [qd dqd] = quinticPoly(t0,q0,dq0,d2q0,tf,qf,dqf,d2qf)
  n = numel(q0);
  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
  qd =@(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)];
  dqd =@(t) (t>=t0 & t<=tf)*dq_poly(t) + (t<t0)*dq0 + (t>tf)*dqf;
end

