[E,t1,q1,q2,qd1,qd2,qdd1,qdd2,torque,g]...

function [E,t1,q1,q2,qd1,qd2,qdd1,qdd2,torque,g]...
    = twolink(Xtraj,Xplant,p)
% Simulate two-link manipulator using quintic trajectory and inverse
% dynamics to compute input torques and energy consumption. Designed for
% traj-only optimization with plant design passed in. No control parameters
% specified since inverse dynamics used.

tf = p.tf;          % final time (assume t0 = 0)
dt = p.dt;          % time step
ploton = p.ploton;  % plotting flag
task = p.task;      % task specification (1 = A, 2 = B)
Tallow = p.Tallow;  % allowable joint torque

% Prepare trajectory parameters:
pI = Xtraj(1:2); vI = Xtraj(3:4);

% Prepare plant design vector: 
X = Xplant(1:2);       % link lengths (m)

% trajactory boundary conditions:
if task == 1
    p0 = [0.1 0.1]';
    v0 = [0 0.1]';
    pf = [0.7 0.1]';
    vf = [0 -0.1]';
elseif task == 2
    p0 = [0.5 1.2]';
    v0 = [-0.1 0]';
    pf = [0.4 2.0]';
    vf = [-0.1 0]';
end

% calculate trajectory:
traj = traj_quin(X,tf,p0,v0,pI,vI,pf,vf,p);
if ~isstruct(traj)  % trajectory was infeasible (singular matrix)
    E = 2000;
    g = [2000 2000]';
    torque = 0;t1=0;q1=0;q2=0;qd1=0;qd2=0;qdd1=0;qdd2=0;
    warning('Trajectory infeasible')
    return
end

% compute trajectory data
t1 = (0:dt:tf)'; nt = length(t1);     
q1 = traj.traj_q1(t1);  
q2 = traj.traj_q2(t1);
qd1 = traj.traj_qd1(t1);
qd2 = traj.traj_qd2(t1);
qdd1 = traj.traj_qdd1(t1);
qdd2 = traj.traj_qdd2(t1);

% compute required input torque using inverse dynamics:
u = zeros(nt,2);
for i=1:length(t1)
    u(i,:) = inv_control([q1(i) q2(i)]',[qd1(i) qd2(i)]'...
        ,[qdd1(i) qdd2(i)]',Xplant)';    % size of u is [nt,2]
end

% calculate energy and power
q_2 = [q1, q2, qd1, qd2]; torque = u;
    % net combined power
P = torque(:,1).*q_2(:,3).^2.*sign(q_2(:,3))...
    + torque(:,2).*q_2(:,4).^2.*sign(q_2(:,4)); 
    % joint 1 power, assuming no regenerative braking
P1 = max(torque(:,1).*q_2(:,3).^2.*sign(q_2(:,3)),0);
    % joint 2 power, assuming no regenerative braking
P2 = max(torque(:,2).*q_2(:,4).^2.*sign(q_2(:,4)),0);
    % combined non-regenerative power
Pc = P1 + P2;   
    % total energy consumption
E = sum(Pc(2:end).*diff(t1));

if E > 1000
    E = 1000;      % penalty for infeasible designs
end

% compute torque constraint
Tmax1 = max(abs(torque(:,1)));
Tmax2 = max(abs(torque(:,2)));
g(1) = Tmax1 - Tallow;
g(2) = Tmax2 - Tallow;

if ploton     
    % plot torque: 
    figure(gcf+1);clf
    plot(t1,torque)
    legend('\tau_1','\tau_2')
    
    % plot torque/speed curves:
    figure(gcf+1);clf
    plot(q_2(:,3),torque(:,1),'ko-'); hold on
    plot(q_2(1,3),torque(1,1),'ko','MarkerSize',10,...
        'MarkerFaceColor',[.49 1 .63]); 
    plot(q_2(:,4),torque(:,2),'ro-');
    plot(q_2(1,4),torque(1,2),'ro','MarkerSize',10,...
        'MarkerFaceColor',[.49 1 .63]);
    legend('q_1','q_1 start','q_2','q_2 start','Location','EastOutside')
    xlabel('speed')
    ylabel('torque')
    
    % plot power
    figure(gcf+1);clf
    plot(t1,P,'k-'); hold on
    plot(t1,P1,'rd-')
    plot(t1,P2,'g*-')
    plot(t1,Pc,'b+-')
    legend('P','P_1','P_2','P_c')
end

function traj = traj_quin(x,tf,p0,v0,pI,vI,pf,vf,p)
% Quintic trajectory generation:
% Creates trajectory functions of time if trajectory is feasible, zeros
% otherwise.

ploton = p.ploton2;

% convert to configuration space:
q0 = inv_kin(p0,x);
qi = inv_kin(pI,x);
qdi = inv_vel(vI,pI,x);
if qdi == 0  % singular matrix detected
    traj = 0;
    return
end
qf = inv_kin(pf,x);
qd0 = inv_vel(v0,q0,x);
if qd0 == 0 % singular matrix detected
    traj = 0;
    return
end
qdf = inv_vel(vf,qf,x);
if qdf == 0 % singular matrix detected
    traj = 0;
    return
end

% solve for q1 trajectories
tI = tf/2; 
A1 = [1 0 0 0 0 0
    0 1 0 0 0 0
    1 tI tI^2 tI^3 tI^4 tI^5
    0 1 2*tI 3*tI^2 4*tI^3 5*tI^4 
    1 tf tf^2 tf^3 tf^4 tf^5
    0 1 2*tf 3*tf^2 4*tf^3 5*tf^4];
b1 = [q0(1) qd0(1) qi(1) qdi(1) qf(1) qdf(1)]';
a1 = A1\b1;    

% solve for q2 trajectories
b2 = [q0(2) qd0(2) qi(2) qdi(2) qf(2) qdf(2)]';
a2 = A1\b2;

% form trajectory functions
traj_q1 = @(t) a1(1)+a1(2)*t+a1(3)*t.^2+a1(4)*t.^3+a1(5)*t.^4+a1(6)*t.^5;
traj_q2 = @(t) a2(1)+a2(2)*t+a2(3)*t.^2+a2(4)*t.^3+a2(5)*t.^4+a2(6)*t.^5;
traj_qd1 = @(t) a1(2)+2*a1(3)*t+3*a1(4)*t.^2+4*a1(5)*t.^3+5*a1(6)*t.^4;
traj_qd2 = @(t) a2(2)+2*a2(3)*t+3*a2(4)*t.^2+4*a2(5)*t.^3+5*a2(6)*t.^4;
traj_qdd1 = @(t) 2*a1(3)+6*a1(4)*t+12*a1(5)*t.^2+20*a1(6)*t.^3;
traj_qdd2 = @(t) 2*a2(3)+6*a2(4)*t+12*a2(5)*t.^2+20*a2(6)*t.^3;

% form output argument
traj.traj_q1 = traj_q1;
traj.traj_q2 = traj_q2;
traj.traj_qd1 = traj_qd1;
traj.traj_qd2 = traj_qd2;
traj.traj_qdd1 = traj_qdd1;
traj.traj_qdd2 = traj_qdd2;

% plot
if ploton
    figure(1);clf
    fplot(traj_q1,[0 tf]); hold on
    set(findobj(gca, 'Type', 'Line'), 'LineWidth', 2,'Color', 'k');
    fplot(traj_q2,[0 tf],'r-'); hold on
    set(findobj(gca, 'Type', 'Line'), 'LineWidth', 2);
    legend('q_1','q_2')
    title('Position Trajectories')
    
    % plot
    figure(2);clf
    fplot(traj_qd1,[0 tf]); hold on
    set(findobj(gca, 'Type', 'Line'), 'LineWidth', 2,'Color', 'k');
    fplot(traj_qd2,[0 tf],'r-'); hold on
    set(findobj(gca, 'Type', 'Line'), 'LineWidth', 2);
    legend('qd_1','qd_2')
    title('Velocity Trajectories')
end


function u = inv_control(q,qd,qdd,X)
% Compute required input torque as a function of desired trajectory and the
% extended plant design vector.

% check vector orientation:
if isrow(q)
    q = q';
end
if isrow(qd)
    qd = qd';
end
if isrow(qdd)
    qdd = qdd';
end

q1 = q(1);
q2 = q(2);
q1d = qd(1);
q2d = qd(2);

% physical constants:
g = 9.81;           % gravitational constant

% manipulator parameters:           
l1 = X(1);          % link 1 length
l2 = X(2);          % link 2 length
m1 = X(3);          % link 1 mass
m2 = X(4);          % link 2 mass
mp = X(5);          % payload mass

% dependent parameters:
lc1 = l1/2;         % link 1 cm location
lc2 = l2*(m2/2+mp)/(m2+mp);         
                    % link 2 cm location
I1 = m1*l1^2/12;    % link 1 mass moment of inertia
Ilink = m2*l2^2/12 + m2*(lc2-l2/2)^2;
Ipayload = mp*(l2-lc2)^2;
I2 = Ilink+Ipayload;% link 2 mass moment of inertia

% update total link masses for D(q) calculation:
m2 = m2 + mp;

% form inertia matrix:
D(1,1) = m1*lc1^2 + m2*(l1^2+lc2^2+2*l1*lc2*cos(q2)) + I1 + I2;
D(1,2) = m2*(lc2^2+l1*lc2*cos(q2)) + I2;
D(2,1) = D(1,2);
D(2,2) = m2*lc2^2 + I2;

% form damping matrix:
h = -m2*l1*lc2*sin(q2);
C(1,1) = h*q2d;
C(1,2) = h*q2d + h*q1d;
C(2,1) = -h*q1d;

% form gravity vector:
G(1,1) = (m1*lc1+m2*l1)*g*cos(q1) + m2*lc2*g*cos(q1+q2);
G(2,1) = m2*lc2*g*cos(q1+q2);

% compute required torque (outputs column vector)
u = D*qdd + C*qd + G; 

function qd = inv_vel(xi,q,x)
% Given a velocity of the end effector in the workspace, calculate the
% velocity in the configuration space.

% configuration space coordinates:
q1 = q(1);              % theta 1
q2 = q(2);              % theta 2

% manipulator parameters:
l1 = x(1);              % link 1 length
l2 = x(2);              % link 2 length

% Jacobian:
J = [(-l1*sin(q1) - l2*sin(q1+q2))  (-l2*sin(q1+q2))
     (l1*cos(q1) + l2*cos(q1+q2))   (l2*cos(q1+q2))];

% Configuration space velocity:
if cond(J) > 1000
    qd = 0;
    return
end
qd = J\xi;

function q = inv_kin(X,x)
% Given a position of the end effector in the workspace, calculate the
% position in the configuration space.

% workspace coordinates:
xc = X(1);              % horizontal 
yc = X(2);              % vertical

% manipulator parameters:
l1 = x(1);              % link 1 length
l2 = x(2);              % link 2 length

D = (xc^2+yc^2-l1^2-l2^2)/(2*l1*l2);

% this assumes elbow down position
q2 = atan2(-sqrt(1-D^2),D);
q1 = atan2(yc,xc) - atan2(l2*sin(q2),l1+l2*cos(q2));

q = [q1 q2]';