function IndirectSingleShooting
%% Optimal Control (EE-715)
% Course at EPFL
% Example - Simple indirect single shooting
%
%
% Autor:    timm.faulwasser@kit.edu
% Date:     01-08-16
% Status:   final
 
close all
clc

    %% parameters
    t0 = 0;
    t1 = 5;

    % initial guesses
    nu0 = 0;
    lambda0 = -1;
    
    % initial and boundary conditions
    x0 =  0;
    x1 =  .50;
    dlambda_dlambda0_0 = 1;
    dx_dlambda0_0      = 0;

    max_iter = 100;
    tolerance = 1E-4;
    odeopt = odeset('RelTol', 1e-12, 'AbsTol', 1e-9, 'Jacobian', @jacobian);
    
    handle_fig = figure(1);
    plot_flag = 1;
    
    %% main loop    
    nu = nu0;
    lambda = lambda0;
    
    tic;
    for i = 1:max_iter

        % Integrate ODE
        [t,y] = ode113(@dynamics, [t0, t1], ...
            [x0, lambda, dlambda_dlambda0_0, dx_dlambda0_0]', odeopt);          
        
        %% minimize defect
        % compute defect
        F = defect(y(end, 2), nu, y(end, 1), x1);
        
        % break condition
        if norm(F) <= tolerance
            break
            
        elseif i == max_iter
            warning('no convergence')
            F
        end
        
        % gradient update
            gradF = [y(end, 3), -1; 
                     y(end, 4),  0];
        if rank(gradF) == 2    
            d = linsolve(gradF, -F);
        else keyboard
        end 
        
        % update initial conditions 
        
        lambda  = lambda + d(1);
        nu      = nu     + d(2);
        
        %% plot progress
        u_opt = y(:,2).*(y(:,1)-1); 
        if  plot_flag == 1
            figure(handle_fig);
            subplot(3,1,1)
            plot(t, y(:,1), 'r', 'LineWidth', 2)
            hold on
            
            if i == 2
                grid on
                legend('x')
            end
            
            figure(handle_fig);
            subplot(3,1,2)
            plot(t, y(:,2), 'b', 'LineWidth', 2)
            hold on
            
            if i == 2
                grid on
                legend('\lambda')
            end
            
            figure(handle_fig);
            subplot(3,1,3)
            plot(t, u_opt , 'g', 'LineWidth', 2)
            hold on 
            
            if i == 2
                grid on
                legend('u_{opt}')
            end
            
            plot_flag = 0;
        else
            plot_flag =1;
        end
    end
    t_comp =toc;
    
    display(['Computation time:  ', num2str(t_comp)]);
    display(['Iterations:  ', num2str(i)]);
    display(['Defect:      ', num2str(F(1)),', ',num2str(F(2))]);
    
    %% plotting
        u_opt = y(:,2).*(y(:,1)-1);

        figure(2)
        subplot(3,1,1)
        plot(t, y(:,1), 'r', 'LineWidth', 2)
        grid on
        legend('x')
        subplot(3,1,2)
        plot(t, y(:,2), 'b', 'LineWidth', 2)
        grid on
        legend('\lambda')
        subplot(3,1,3)
        plot(t, u_opt , 'g', 'LineWidth', 2)
        grid on
        legend('u_{opt}')
 %   keyboard
 %   keyboard
end
%% subfunctions
function dx = dynamics(t,x)
    x1                  = x(1);
    lambda              = x(2);
    dlambda_dlambda0    = x(3);
    dx_dlambda0         = x(4);

    u = (x1-1)*lambda;

    dx1                 =  u*(1-x1);
    dlambda             =  u*lambda;
    ddlamdba_dlambda0   =  2*lambda*(x1-1);
    ddx_dlambda0        = -lambda*(2*x1-2);    
    
    dx = [dx1, dlambda, ddlamdba_dlambda0, ddx_dlambda0]';
end

function J = jacobian(t,x)
    x1                 = x(1);
    lambda             = x(2);
    dlambda_dlambda0   = x(3);
    dx_dlambda0        = x(4);

    J      =   zeros(4,4);
    J(1,1) =   2*lambda*(x1-1);
    J(1,2) =  -(x1-1)^2;
    J(2,1) =   lambda^2;
    J(2,2) =   2*lambda*(x1-1);
    J(3,1) =   2*lambda;
    J(3,2) =   lambda;
    J(4,1) =  -2*lambda;
    J(4,2) =   2*x1-2;
end


function F = defect(lambda, nu, x, x1)   
    F1 = lambda - nu;
    F2 = terminal_constraint(x, x1);
    F  = [F1; F2];
end   

function psi = terminal_constraint(x, x1)
    psi = x - x1;
end
