function [ tsol, xsol, usol, Isol ] = robopt ( N, nit, c, A, B );
% RobOpt
%
% Optimal trajectory of a simplified 2D Robot model
%
% This m-file is part of the working materials of the 
%   Schuelerseminar "Optimale Bahnplanung von Robotern"
% at the Department of Mathematics and Computer Science at
% Martin-Luther University Halle-Wittenberg, Germany
%
% Parameters
%   N      (input)  : number of discretization points
%   nit    (input)  : number of iteration steps in Newton's method
%   c      (input)  : [ c1  c2  ], weighting factors in functional
%   A      (input)  : [ x_A y_A ], initial position
%   B      (input)  : [ x_B y_B ], final position
%   tsol   (output) : time grid
%   xsol   (output) : numerical solution [ x1(t) x2(t) x3(t) x4(t) ]
%   usol   (output) : optimal control    [ u1(t) u2(t) ]
%   Isol   (output) : functional \int_0^t c_1u_1^2+c_2u_2^2 dt
%
% Example
%   [ tsol, xsol, usol, Isol ] = robopt ( 200, 10, [ 1 1 ], [ 0.8 0.2 ], [ -0.8 0.2 ] );
%
% Author :      Prof. Dr. M. Arnold, martin.arnold@mathematik.uni-halle.de
% Version of :  Feb 11, 2006

% -> problem data
t_f = 1.0;   % [s]       ... length of time interval

len = 0.5;   % [m]       ... length of robot arm
ma  = 1.0;   % [kg]      ... mass of robot arm
ja  = 1.0;   % [kg m^2]  ... moment of inertia of robot arm
ml  = 0.5;   % [kg]      ... mass of payload
jl  = 1.0;   % [kg m^2]  ... moment of inertia of payload

robdat = [ t_f, len, ma, ja, ml, jl ];

% -> read input
x_A = A(1);
y_A = A(2);
x_B = B(1);
y_B = B(2);

% -> initial guess
x1_0 = atan ( y_A / x_A );   if x1_0<0,  x1_0 = x1_0 + pi;  end;
x2_0 = sqrt ( x_A^2 + y_A^2 );
x1_f = atan ( y_B / x_B );   if x1_f<0,  x1_f = x1_f + pi;  end;
x2_f = sqrt ( x_B^2 + y_B^2 );

tsol = linspace(0,t_f,N+1);
xsol = [ linspace(x1_0,x1_f,N+1)'  ...
       , linspace(x2_0,x2_f,N+1)'  ...
       , zeros(N+1,1)              ...
       , zeros(N+1,1)              ...
       ];
lsol  = zeros ( N+1, 4 );
[ usol, Isol ]  = robgetu ( xsol, lsol, c, robdat );

robplot ( tsol, xsol, usol, Isol, NaN, 0, NaN, robdat, A, B );

% -> Newton iteration to get numerical solution by trapezoidal rule
for iit=1:nit,

    % -> get function value
    F = robf ( xsol, lsol, c, robdat );

    % -> difference approximation of Jacobian
    DF = zeros ( 8*N, 8*N ); 
    x = robsol2x ( xsol, lsol );

    for i=1:(8*N),
        dx    = x;
        del   = sqrt(eps) * max ( abs(dx(i)), 1.0e-4 );
        dx(i) = dx(i) + del;
        [ xdel, ldel ] = robx2sol ( dx, xsol(1+0,:), xsol(1+N,:) );
        Fdel = robf ( xdel, ldel, c, robdat );
        DF(:,i) = ( Fdel - F ) / del;
    end;

    % -> Newton step
    p = - DF \ F;
    x = robsol2x ( xsol, lsol );
    x = x + p;
    [ xsol, lsol ] = robx2sol ( x, xsol(1+0,:), xsol(1+N,:) );

    % -> visualization
    [ usol, Isol ]  = robgetu ( xsol, lsol, c, robdat );
    robplot ( tsol, xsol, usol, Isol, 4, iit, norm(F), robdat, A, B );

end;

%---------------------------------------------------------%

function F = robf ( xsol, lsol, c, robdat );
% RobF
%
% Funtion value in trapezoidal rule
%
% Parameters
%   xsol   (input)  : numerical solution [ x1(t) x2(t) x3(t) x4(t) ]
%   lsol   (input)  : Lagrangian multipliers
%   c      (input)  : [ c1  c2  ], weighting factors in functional
%   robdat (input)  : physical parameters of robot model
%   F      (output) : Function value for Newton's method

% -> read input 
t_f = robdat(1);
len = robdat(2);
ma  = robdat(3);
ja  = robdat(4);
ml  = robdat(5);
jl  = robdat(6);

% -> allocate memory 
N = length ( xsol(:,1) ) - 1;
F = zeros ( 8*N, 1 );

% -> get current control
[ usol, Isol ]  = robgetu ( xsol, lsol, c, robdat );

% -> evaluate right hand side
frhs = zeros ( N+1, 8 );

frhs(:,1) = xsol(:,3);
frhs(:,2) = xsol(:,4);
frhs(:,3) = usol(:,1) ./ ( ma*(xsol(:,2)-len).^2 + ja + ml*xsol(:,2).^2 + jl );
frhs(:,4) = ( usol(:,2) + ( ma*(xsol(:,2)-len) + ml*xsol(:,2) ) .* xsol(:,3).^2 ) / ( ma + ml );
frhs(:,5) = zeros ( N+1, 1 );
frhs(:,6) = lsol(:,3) .* usol(:,1) .* ( 2*ma*(xsol(2)-len) + 2*ml*xsol(:,2) )  ...
              ./ ( ma*(xsol(2)-len).^2 + ja + ml*xsol(2).^2 + jl ).^2          ...
              - lsol(:,4) .* xsol(:,3).^2;
frhs(:,7) = - lsol(:,1) - lsol(:,4) .* ( 2*ma*(xsol(2)-len) + 2*ml*xsol(:,2) ) .* xsol(:,3);
frhs(:,8) = - lsol(:,2);

% -> discretized time derivatives
h = 1/N;
for j=1:4,
    F(    (j-1)*N+(1:N)) = diff(xsol(:,j)) / h;
    F(4*N+(j-1)*N+(1:N)) = diff(lsol(:,j)) / h;
end;

% -> right hand sides
for j=1:8,
    F((j-1)*N+(1:N)) = F((j-1)*N+(1:N)) - ( frhs(1+(0:(N-1)),j) + frhs(1+(1:N),j) ) / 2;
end;

%---------------------------------------------------------%

function x = robsol2x ( xsol, lsol );
% RobSol2x
%
% Save solution [ xsol, lsol ] in vector x
%
% Parameters
%   xsol   (input)  : numerical solution [ x1(t) x2(t) x3(t) x4(t) ]
%   lsol   (input)  : Lagrangian multipliers
%   x      (output) : solution vector for Newton's method

N = length ( xsol(:,1) ) - 1;
x = [ xsol(1+(1:N-1),1); xsol(1+(1:N-1),2)       ...
    ; xsol(1+(1:N-1),3); xsol(1+(1:N-1),4)       ...
    ; lsol(:,1); lsol(:,2); lsol(:,3); lsol(:,4) ];

%---------------------------------------------------------%

function [ xsol, lsol ] = robx2sol ( x, xsol_0, xsol_f );
% Robx2Sol
%
% Get solution [ xsol, lsol ] from vector x
%
% Parameters
%   x      (input)  : solution vector from Newton's method
%   xsol_0 (input)  : numerical solution x(0)
%   xsol_f (input)  : numerical solution x(t_f)
%   xsol   (output) : numerical solution [ x1(t) x2(t) x3(t) x4(t) ]
%   lsol   (output) : Lagrangian multipliers

N = length(x) / 8;

xsol = zeros ( N+1, 4 ); 
lsol = zeros ( N+1, 4 );

xsol(1+0,:) = xsol_0;
xsol(1+N,:) = xsol_f;

for j=1:4,
    xsol(1+(1:N-1),j) = x((j-1)*(N-1)+(1:(N-1)));
    lsol(    :    ,j) = x(4*(N-1)+(j-1)*(N+1)+(1:(N+1)));
end;

%---------------------------------------------------------%

function [ usol, Isol ] = robgetu ( xsol, lsol, c, robdat );
% RobGetU
%
% Evaluate optimal control u(t)
%
% Parameters
%   xsol   (input)  : numerical solution [ x1(t) x2(t) x3(t) x4(t) ]
%   lsol   (input)  : Lagrangian multipliers
%   c      (input)  : [ c1  c2  ], weighting factors in functional
%   robdat (input)  : physical parameters of robot model
%   usol   (output) : current control [ u1(t) u2(t) ]
%   Isol   (output) : functional \int_0^t c_1u_1^2+c_2u_2^2 dt

% -> read input 
t_f = robdat(1);
len = robdat(2);
ma  = robdat(3);
ja  = robdat(4);
ml  = robdat(5);
jl  = robdat(6);

% -> get control from Euler-Lagrange equations
N = length ( xsol(:,1) ) - 1;
h = t_f / N;

usol = zeros ( N+1, 2 );

usol(:,1) = - lsol(:,3) ./ ( 2*c(1)*(ma*(xsol(2)-len).^2+ja+ml*xsol(2).^2+jl) );
usol(:,2) = - lsol(:,4) / ( 2*c(2)*(ma+ml) );

% -> evaluate functional by trapezoidal rule
Isol = zeros ( N+1, 1 );
Isol(1+0) = h/2 * ( c(1) * usol(1+0,1)^2 + c(2) * usol(1+0,2)^2 );
for i=1:N-1,
    Isol(1+i) = Isol(1+i-1) + h * ( c(1) * usol(1+i,1)^2 + c(2) * usol(1+i,2)^2 );
end;
Isol(1+N) = Isol(1+N-1) + h/2 * ( c(1) * usol(1+N,1)^2 + c(2) * usol(1+N,2)^2 );

%---------------------------------------------------------%

function robplot ( tsol, xsol, usol, Isol, fani, iit, normF, robdat, A, B );
% RobPlot
%
% Plot of results
%
% Parameters
%   tsol   (input)  : time grid
%   xsol   (input)  : numerical solution [ x1(t) x2(t) x3(t) x4(t) ]
%   usol   (input)  : current control [ u1(t) u2(t) ]
%   Isol   (input)  : functional \int_0^t c_1u_1^2+c_2u_2^2 dt
%   fani   (input)  : animation frequency, set fani=NaN to skip animation
%   iit    (input)  : number of iteration step
%   normF  (input)  : norm of residual
%   robdat (input)  : physical parameters of robot model
%   A      (input)  : [ x_A y_A ], initial position
%   B      (input)  : [ x_B y_B ], final position

% -> read input
len = robdat(2);

% -> plot trajectory
subplot ( 2, 2, 1 );
h1_blue = plot (                             ...
                 xsol(:,2).*cos(xsol(:,1))   ...
               , xsol(:,2).*sin(xsol(:,1))   ...
               , 'b'                         ...
               , 'LineWidth', 2              ...
               );
hold on
h1_cyan = plot (                             ...
                 xsol(1,2).*cos(xsol(1,1))   ...
               , xsol(1,2).*sin(xsol(1,1))   ...
               , '.c'                        ...
               , 'MarkerSize', 16            ...
               );
h1_black = plot ( (xsol(1,2)-len*[2 0]).*cos(xsol(1,1))   ...
                , (xsol(1,2)-len*[2 0]).*sin(xsol(1,1))   ...
                , '-w'                                    ...
                , 'LineWidth', 1                          ...
                );

plot (                             ...
       linspace(xsol(1,2),xsol(end,2),201).*cos(linspace(xsol(1,1),xsol(end,1),201))  ...
     , linspace(xsol(1,2),xsol(end,2),201).*sin(linspace(xsol(1,1),xsol(end,1),201))  ...
     , ':k'                                                                           ...
     );
plot ( 0,    0,    '+k', 'MarkerSize', 12 );
plot ( A(1), A(2), '.r', 'MarkerSize', 24 );
plot ( B(1), B(2), '.r', 'MarkerSize', 24 );
hold off
xlabel ( 'x' );
ylabel ( 'y' );
title ( sprintf('Payload trajectory: Iterationsschritt %2i,  Residuum %6.2e',iit,normF) )
axis ( [ -1 1 -1 1 ] );
set ( gca, 'PlotBoxAspectRatio', [1 1 1] );
set ( gca, 'DataAspectRatio',    [1 1 1] );

% -> plot solution
subplot ( 2, 2, 2 );
h2 = plot ( tsol, xsol(:,1:2) );
xlabel ( 't' );
ylabel ( 'x(t)' );
title ( 'Numerical solution' )
legend ( 'x_1(t)', 'x_2(t)'                     ...
       , 'Location', 'NorthWest'                ...
       );

% -> plot control
subplot ( 2, 2, 3 );
h3 = plot ( tsol, usol(:,1:2) );
xlabel ( 't' );
ylabel ( 'u(t)' );
title ( 'control' )
legend ( 'u_1(t)', 'u_2(t)'                     ...
       , 'Location', 'NorthEast'                ...
       );

% -> plot functional
subplot ( 2, 2, 4 );
h4 = plot ( tsol, Isol );
xlabel ( 't' );
ylabel ( 'I(t)' );
title ( 'Energy consumption' )
legend ( '\int_0^t c_1u_1^2(t)+c_2u_2^2(t) dt'  ...
       , 'Location', 'NorthWest'                ...
       );

% -> plot
drawnow

% -> animation
if ~isnan(fani),

    pause ( 2 );
    
    N = length ( xsol(:,1) ) - 1;
    set ( h1_black, 'Color', 'k' );
    for ax=get(gcf,'Children')',
        set ( ax, 'XLimMode', 'manual' );
        set ( ax, 'YLimMode', 'manual' );
    end;

    for i=0:fani:N,
        set ( h1_blue,  'XData', xsol(1+(0:i),2).*cos(xsol(1+(0:i),1)) );
        set ( h1_blue,  'YData', xsol(1+(0:i),2).*sin(xsol(1+(0:i),1)) );
        set ( h1_cyan,  'XData', xsol(1+i,    2).*cos(xsol(1+i,    1)) );
        set ( h1_cyan,  'YData', xsol(1+i,    2).*sin(xsol(1+i,    1)) );
        set ( h1_black, 'XData', (xsol(1+i,2)-len*[2 0]).*cos(xsol(1+i,1)) );
        set ( h1_black, 'YData', (xsol(1+i,2)-len*[2 0]).*sin(xsol(1+i,1)) );
        set ( h2(1), 'XData', tsol(1+(0:i)) );
        set ( h2(1), 'YData', xsol(1+(0:i),1) );
        set ( h2(2), 'XData', tsol(1+(0:i)) );
        set ( h2(2), 'YData', xsol(1+(0:i),2) );
        set ( h3(1), 'XData', tsol(1+(0:i)) );
        set ( h3(1), 'YData', usol(1+(0:i),1) );
        set ( h3(2), 'XData', tsol(1+(0:i)) );
        set ( h3(2), 'YData', usol(1+(0:i),2) );
        set ( h4,    'XData', tsol(1+(0:i)) );
        set ( h4,    'YData', Isol(1+(0:i),1) );
        drawnow
    end;

end;
