[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
% since inverse dynamics used.
%
% Author: James T. Allison, Assistant Professor, University of Illinois at
% Urbana-Champaign
% Date: 7/30/12
% Copyright 2012, James T. Allison

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);

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

% 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:
for i=1:length(t1)
    u(i,:) = inv_control([q1(i) q2(i)]',[qd1(i) qd2(i)]'...
        ,[qdd1(i) qdd2(i)]',X)';    % 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);

% calculate inertia properties
[m1, m2, lc1, lc2, I1, I2] = calc_I(X);

% link lengths
l1 = X(1);          % link 1 length
l2 = X(2);          % link 2 length

% calculate dynamic model matrices
[D,C,G] = calc_D(m1, m2, lc1, lc2, l1, l2, I1, I2, q1, q2, q1d, q2d);

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

function [m1, m2, lc1, lc2, I1, I2] = calc_I(X)
% Calculate inertia properties for counter-balanced manipulator.

% 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
lcw1 = X(6);        % counterweight 1 link length
lcw2 = X(7);        % counterweight 2 link length
mlc1 = X(8);        % counterweight link 1 mass
mlc2 = X(9);        % counterweight link 2 mass
mc1 = X(10);        % counterweight 1 mass
mc2 = X(11);        % counterweight 2 mass

% center of mass locations:
lc1 = (0.5*l1*m1 - 0.5*lcw1*mlc1 - lcw1*mc1)/...
    (m1 + mlc1 + mc1);      % link 1 cm location
lc2 = (0.5*l2*m2 - 0.5*lcw2*mlc2 - lcw2*mc2 + l2*mp)/...
    (m2 + mlc2 + mc2 + mp); % link 2 cm location

% moments of inertia
    % I link 1 about new cm
I1n = m1*l1^2/12 + m1*(0.5*l1 - lc1)^2;
    % I counterweight link 1 about new cm
Icw1n = mlc1*lcw1^2/12 + mlc1*(0.5*lcw1 + lc1)^2;
    % I counterweight 1 about new cm
Imc1 = mc1*(lcw1 + lc1)^2;
    % total link 1 I about new cm:
I1 = I1n + Icw1n + Imc1;

    % I link 2 about new cm
I2n = m2*l2^2/12 + m2*(0.5*l2 - lc2)^2;
    % I counterweight link 1 about new cm
Icw2n = mlc2*lcw2^2/12 + mlc2*(0.5*lcw2 + lc2)^2;
    % I counterweight 1 about new cm
Imc2 = mc2*(lcw2 + lc2)^2;
    % I payload about new cm
Ipn = mp*(lcw2 - lc2)^2;
    % total link 1 I about new cm:
I2 = I2n + Icw2n + Imc2 + Ipn;

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

function [D,C,G] = calc_D(m1,m2,lc1,lc2,l1,l2,I1,I2,q1,q2,q1d,q2d)
% Calculate dynamic model matrices for counter-balanced manipulator.

% physical constants:
g = 9.81;           % gravitational constant

% 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);

function [Xplant,r1,r2] = calc_r(X)

% independent plant design variables:
l1 = X(1);
l2 = X(2);
lcw1 = X(3);
lcw2 = X(4);
mc1 = X(5);
mc2 = X(6);
% Sample design: X = [1.0 1.0 0.3 0.3 10 10]';

% plant design parameters:
mp = 20;            % payload mass
E = 71.7e9;         % Elastic modulus, Pa
delta_allow = 1e-5; % allowable deflection (should I change to deflection/length)?
t1 = 0.002;         % link 1 wall thickness (m)
t2 = 0.002;         % link 2 wall thickness (m)
rho = 2810;         % Material density (kg/m^3) 7075 Aluminum
tau1_max = 350;     % Joint 1 nominal torque
tau2_max = 350;     % Joint 2 nominal torque

% intermediate plant calculations:

% link 1:
b1 = max(lcw1,l1);      % greatest distance from joint to end
a1 = lcw1;              % counterweight link length
lt1 = lcw1 + l1;        % total link length
C1 = asin(delta_allow/b1);   % intermediate variable
D1 = (tau1_max/(6*E))*(6*a1^2-6*a1*lt1+2*lt1^2);
G1 = 4*D1/(pi*C1);
r1 = (G1+t1^2)/(2*t1);
m1 = l1*rho*pi*(r1^2-(r1-t1)^2);
mlc1 = lcw1*rho*pi*(r1^2-(r1-t1)^2);

% link 1:
b2 = max(lcw2,l2);      % greatest distance from joint to end
a2 = lcw2;              % counterweight link length
lt2 = lcw2 + l2;        % total link length
C2 = asin(delta_allow/b2);   % intermediate variable
D2 = (tau2_max/(6*E))*(6*a2^2-6*a2*lt2+2*lt2^2);
G2 = 4*D2/(pi*C2);
r2 = (G2+t2^2)/(2*t2);
m2 = l2*rho*pi*(r2^2-(r2-t2)^2);
mlc2 = lcw2*rho*pi*(r2^2-(r2-t2)^2);

% complete plant parameter vector:
Xplant = [l1 l2 m1 m2 mp lcw1 lcw2 mlc1 mlc2 mc1 mc2]';

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]';