function kerbal_solution()
    
    global N iter
    
    iter = 1;
    N=40;
    
    
    x0 = [50;0;10];
    u0 = 0;
    h0 = 4;
    
    
    ny = 3*(N+1)+N+1;
    
    y0 = x0;
    y0 = [y0; repmat([u0;x0],N,1)];
    y0 = [y0;h0];
    
    lb = -inf*ones(ny,1);
    ub = inf*ones(ny,1);
    
    % x0 constraint
    lb(1:3) = x0;   ub(1:3) = x0;           % x(0)   = x0
    
    
    % final constraints
    lb(4*N+1) = 0;    ub(4*N+1) = 0;        % p(T)   =  0
    lb(4*N+2) = 0;    ub(4*N+2) = 0;        % v(T)   =  0
    lb(4*N+3) = 4;                          % mf(T) >=  4
    
    % positive time
    lb(end) = 0;    % ub(end) = 4.8;
    
    % make y0 satisfy constraints
    y0(4*N) = 0;
    y0(4*N+1) = 0;
    
    
    opts = optimset('GradObj','on','GradConstr','on','display','final','algorithm','sqp',...
        'MaxFunEvals',1000000,'OutputFcn',@(x,y,z)plot_iter(x,y,z,N));
    tic
    [y_opt,fval,exitflag,output,lambda,grad,hess] = fmincon(@objective, y0, [], [], [], [], lb, ub, @nonlin_constraints, opts);
    toc

end





function stop = plot_iter(x,optimValues,state,N)

    stop = false;
    switch state
        case 'init'
            figure(1)
        case 'iter'
            % Concatenate current point and objective function
            % value with history. x must be a row vector.
            clf
            plot_trajectory(x,N)

            drawnow
            

        case 'done'

        otherwise
    end
end


function plot_trajectory(y,N)

        global iter


        p_opt = y(1:4:end-1);
        v_opt = y(2:4:end-1);
        m_opt = y(3:4:end-1);
        u_opt = y(4:4:end-1);
        h_opt = y(end);

        times = 0:h_opt:N*h_opt;
        subplot(2,2,1)
        plot(times,p_opt)

        ylabel('p*');
        xlabel('time');

        subplot(2,2,2)
        plot(times,v_opt)

        ylabel('v*');
        xlabel('time');

        subplot(2,2,3)
        plot(times,m_opt)

        ylabel('mF*');
        xlabel('time');

        subplot(2,2,4)
        stairs(times(1:end-1),u_opt)

        ylabel('u*');
        xlabel('time');

        fprintf('optimal time: %.2f, iteration: %i\n', N*h_opt,iter);
        
        iter=iter+1;


end



function [c,ceq,gradc,gradceq] = nonlin_constraints(y)

    global N

    c = [];
    
    ceq = []; 
    for i=1:N
        
        xi = xi_func(y,i);
        ui = ui_func(y,i);
        xiplus = xi_func(y,i+1);
        
        h = y(end);
        
        ceq = [ceq;integrate_rk4(xi,ui,h)-xiplus];
    end
    
    if nargout > 2
        gradc = [];
        gradceq = GJacFast(y);
    end
end

function [jac] = GJacFast(y)

    global N
    
    h = y(end);

        
    jac = zeros(3*N,numel(y));
    
    
    for i=1:N
        
        row = (i-1)*3 + 1;
        col = (i-1)*4 + 1;
        
        xi = xi_func(y,i);
        ui = ui_func(y,i);
        
        [A,B,X] = rk4step_jac(xi,ui,h);
    
        
        jac(row:row+2,col:col+2) = A;
        jac(row:row+2,col+3) = B;
        jac(row:row+2,col+4:col+6) = -eye(3,3);
        
        jac(row:row+2,end) = X;
        
    end
    
    jac = jac';
    

end


function [A,B,X] = rk4step_jac(x,u,h)
    
    global delta
    
    % Task 3: ...
    
    delta = 1e-4;

    xnext = integrate_rk4(x,u,h);
    
    A(:,1) = (integrate_rk4(x+[delta;0;0],u,h) - xnext) / delta;
    A(:,2) = (integrate_rk4(x+[0;delta;0],u,h) - xnext) / delta;
    A(:,3) = (integrate_rk4(x+[0;0;delta],u,h) - xnext) / delta;

    B = (integrate_rk4(x,u+delta,h) - xnext) / delta;
    
    X = (integrate_rk4(x,u,h+delta) - xnext) / delta;

end



function xi = xi_func(y,i)

    xi = y(4*(i-1)+1:4*(i-1)+3);

end

function ui = ui_func(y,i)

    ui = y(4*(i-1)+4);

end


function [f_obj,grad_obj] = objective(y)

    h = y(end);
    f_obj = h;
    
    grad_obj = zeros(numel(y),1);
    grad_obj(end) = 1;

end


function [ x_next] = integrate_rk4(x,u,h)

    k1 = ode(x,u);
    k2 = ode(x+h/2.*k1,u);
    k3 = ode(x+h/2.*k2,u);
    k4 = ode(x+h*k3,u);
    x_next = x + h/6. * (k1+2*k2+2*k3+k4);
end


function [xdot] =  ode(x,u)

    p=x(1);
    v=x(2);
    mF=x(3);
    
    xdot=[v;u/(30+mF);-u^2];

end