function [tout,finalstate]=planarchain(n,a,b,m,l,g,k,rtol)   %Maciej's ordering of the input.
%planarchain.m solves for the dynamics of a chain of point masses confined
%to a plane.  parameters that can be varied include:
%
%       n is the number of point masses in the chain
%       a is a follower force parameter
%       b is a damping coefficient (damping occurs at each joint)
%       m is the value of each mass
%       l is the distance between adjacent masses
%       g is the value of gravity
%       k is a spring coefficient (each joint is sprung as well)
%
%try: >>planarchain(10,0,0,0.1,0.1,1,0,1e-3);

%establish parameters:
%n=2;%10;
%a=1;%1.0*m*g;  %stability threshhold for these values is about 1.9
%b=0.1;%0.005;
%m=0.2;%1e-2;
%l=0.2;%0.1;
%g=9.8;
%k=0;%0.001;

%establish desired time interval and initial configuration:
time=linspace(0,4,500);
if 0    %chaotic case
thin=zeros(1,n);    %theta values
thdin=zeros(1,n);    %theta derivatives
thdin(n)=5/l;
else     %smooth case
thin=zeros(1,n);    %theta values
thdin=zeros(1,n);    %theta derivatives
thin(1)=pi/4;
end

%solve the ODE.  
%@chain solves the nonlinear and @linchain solves the linear equations.
if 1    %Solve the nonlinear equations
    opt=odeset('RelTol',rtol);
    tic
    [t,state]=ode45(@chain,time,[thin thdin]',opt,n,a,b,m,l,g,k);
    tout=toc; 
    %finalstate: (erase when done with comparing algorithms)
        th=state(end,1:n);
        thd=state(end,n+1:end);
        x=zeros(2,n);
        x(:,1)=l*[cos(th(1));sin(th(1))];
        for j=2:n
            x(:,j)=x(:,j-1)+l*[cos(sum(th(1:j)));sin(sum(th(1:j)))];
        end
        thc=th';        thdc=thd';
        ph=zeros(n,1);  phd=zeros(n,1);
        ph(1)=thc(1);   phd(1)=thdc(1);
        for j=2:n
            ph(j)=ph(j-1)+thc(j);    phd(j)=phd(j-1)+thdc(j);
        end
        sines=sin(ph);  cosines=cos(ph);
        B1=tril(ones(n,1)*cosines');
        B2=tril(ones(n,1)*sines');
        B=[B1(:) B2(:)]';
        B=reshape(B(:),2*n,n);
        A=[-B2(:) B1(:)]';
        A=reshape(A(:),2*n,n)*tril(ones(n));
        xd=A*thdc;
        finalstate=[x reshape(xd,2,n)];
        %return
else    %Solve the linear equations
    opt=odeset('RelTol',1e-4);
    SYS=linear(n,a,b,m,l,g,k);
    [t,state]=ode45(@linchain,time,[thin thdin]',opt,SYS);
end
th=state(:,1:n);
thd=state(:,n+1:end);

%animate!
fig1=figure;  
set(fig1,'color',[1 1 1],'backingstore','off','Doublebuffer','on');
for i=1:length(t)
    %construct a 2xN array of position vectors
    x=zeros(2,n);
    x(:,1)=l*[cos(th(i,1));sin(th(i,1))];
    for j=2:n
        x(:,j)=x(:,j-1)+l*[cos(sum(th(i,1:j)));sin(sum(th(i,1:j)))];
    end
    plot([0 x(1,:)],[0 x(2,:)],'b',x(1,:),x(2,:),'ro',...
         0.5*[-l l],[0 0],'k',[0 0],0.5*[-l l],'k');
    axis equal
    xlim([-l*n,l*n])
    ylim([-l*n,l*n])
    drawnow
end
%title('constraint enforced by the coordinitization')

%Energy Plots:
PE=zeros(length(t),1);
KE=zeros(length(t),1);
for i=1:length(t)
    thc=th(i,:)';    thdc=thd(i,:)';
    ph=zeros(n,1);  phd=zeros(n,1);
    ph(1)=thc(1);   phd(1)=thdc(1);
    for j=2:n
        ph(j)=ph(j-1)+thc(j);    phd(j)=phd(j-1)+thdc(j);
    end
    sines=sin(ph);  cosines=cos(ph);
    B1=tril(ones(n,1)*cosines');
    B2=tril(ones(n,1)*sines');
    B=[B1(:) B2(:)]';
    B=reshape(B(:),2*n,n);
    A=[-B2(:) B1(:)]';
    A=reshape(A(:),2*n,n)*tril(ones(n));
    Q=[zeros(1,n);ones(1,n)];
    PE(i)=m*g*l*(0.5*n*(n+1)-Q(:)'*A(:,1));
    KE(i)=0.5*m*l^2*thdc'*(A'*A)*thdc;
end
fig2=figure;
set(fig2,'color',[1 1 1])
plot(time,KE,'r:',time,PE,'b:',time,KE+PE,'k')
legend('Kinetic Energy','Potential Energy','Total Energy')
xlabel('time')
ylabel('energy')

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function dstate=chain(t,state,n,a,b,m,l,g,k)
%nonlinear equations of motion:
th=state(1:n);  thd=state(n+1:2*n);
ph=zeros(n,1);  phd=zeros(n,1);
ph(1)=th(1);    phd(1)=thd(1);
for i=2:n
    ph(i)=ph(i-1)+th(i);    phd(i)=phd(i-1)+thd(i);
end
sines=sin(ph);  cosines=cos(ph);
B1=tril(ones(n,1)*cosines');
B2=tril(ones(n,1)*sines');
B=[B1(:) B2(:)]';
B=reshape(B(:),2*n,n);
A=[-B2(:) B1(:)]';
A=reshape(A(:),2*n,n)*tril(ones(n));
%Now for the forces:
C=zeros(2*n,n);
C(1:2,1)=[sines(1);-cosines(1)];
C(1:4,2)=[-sines(1)-sines(2);cosines(1)+cosines(2);sines(2);-cosines(2)];
for i=3:n
    C(2*i-5:2*i,i)=[sines(i-1);-cosines(i-1);-sines(i-1)-sines(i);
                    cosines(i-1)+cosines(i);sines(i); -cosines(i)];
end
fg=[ones(1,n);zeros(1,n)];
fg=m*g*fg(:);  
if 0    %follower force is applied only to the last mass
    ff=zeros(2*n,1);
    ff(2*n-1:2*n,1)=-a*[cosines(n);sines(n)];
else    %a follower force is applied to every mass
    ff=[cosines';sines'];
    ff=-a/n*ff(:);  %note that a has been rescaled
end
%Dynamics:
dstate=[thd;
        (A'*A)\(A'*B*phd.^2+(1/m/l)*A'*(ff+fg+C*((b/l)*thd +(k/l)*th)))];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function dstate=linchain(t,state,SYS)  
%linearized equations of motion:
dstate=SYS*state;