%% Gradient method for continuous-time optimal control - fixed final time, free final state

%%
%

x0 = [0.05; 0];                     % the initial state
tspan = linspace(0,0.78,100);       % time span

Sf = diag([1000,1000]);             % coefficient for the quadratic penalty of the final state

udata = ones(length(tspan),1);      % the initial control signal
tu = tspan;                         % and its dedicated time (for generality in future versions)

%% Deriving costate equations from the Hamiltonian

% Building the costate equations (to be solved in the code above)

x = sym('x',[2;1]);
p = sym('p',[2;1]);
u = sym('u');

% definition of the right hand sides of the two state equations

f = [-2*(x(1)+0.25) + (x(2)+0.5) .* exp(25*x(1)./(x(1)+ 2)) - (x(1)+0.25) .* u;
     0.5 - x(2) - (x(2)+0.5) .* exp(25*x(1)./(x(1)+2))];
 
R = 1;                               % the coefficient penalizing the control in the quadratic criterion
L = x(1)^2 + x(2)^2 + R*u^2;         % Lagrangian (the integrand in the cost function)

H = L + p.'*f;                       % Hamiltonian (the control theory format)

% The costate equations. First in the symbolic form and then as an
% anonymous function

dHdx = [diff(H,x(1)); diff(H,x(2))];
dHdx = matlabFunction(dHdx);

% The stationarity equation. First in the symbolic form and then as an
% anonymous function

dHdu = diff(H,u);
dHdu = matlabFunction(dHdu);

%% The iterations

gamma1 = 1e-2;
norm_dHdu = gamma1 + 1;

while norm_dHdu > gamma1
    % Plot u
    figure(1)
    plot(tu,udata);
    xlabel("t");
    ylabel("u(t)");
    
    % Solving the state equations for the given initial state and the given
    % input

    [tx,xdata] = ode45(@(t,x) state_equations(t,x,tu,udata),tspan,x0);

    % Plotting the solution of the state equations

    figure(2)
    subplot(2,1,1);
    plot(tx,xdata(:,1));
    ylabel("x_1(t)");
    subplot(2,1,2);
    plot(tx,xdata(:,2));
    xlabel("t");
    ylabel("x_2(t)");

    % Boundary conditions for the costate equations for the free final state

    xf = xdata(end,:)';
    pf = Sf*xf;                         

    % Solving the costate equations

    [tp,pdata] = ode45(@(t,p) costate_equations(t,p,tx,xdata,tu,udata,dHdx),fliplr(tspan),pf);
    tp = fliplr(tp);
    pdata = fliplr(pdata);
    
    % Plotting the solution of the costate equations

    figure(3)
    subplot(2,1,1);
    plot(tp,pdata(:,1));
    ylabel("p_1(t)");
    subplot(2,1,2);
    plot(tp,pdata(:,2));
    xlabel("t");
    ylabel("p_2(t)");

    %% Criterion for finishing
    norm_dHdu = norm(dHdu(pdata(:,1),udata,xdata(:,1)));
    
    %% Steepest descend by modifying u

    tau = 0.02;                          % this seems to be the most critical parameter (as usual in the gradient method)
    dudata = -tau*dHdu(pdata(:,1),udata,xdata(:,1));
    udata = udata + dudata;
    
    %% Checking the final status
    figure(4)
    plot(tu,dHdu(pdata(:,1),udata,xdata(:,1)));
    xlabel("t");
    ylabel("\partial H / \partial u");
    
end


%% Definition of the (right hand side of the) state equations for the simulation

function dxdt = state_equations(t,x,tu,udata)
ut = interp1(tu,udata,t);
dxdt = zeros(2,1);
dxdt(1) = -2*(x(1)+0.25) + (x(2)+0.5) .* exp(25*x(1)./(x(1)+ 2)) - (x(1)+0.25) .* ut;
dxdt(2) = 0.5 - x(2) - (x(2)+0.5) .* exp(25*x(1)./(x(1)+2));
end

%% Definition of the (right hand side of the) costate equations for the simulation

function dpdt = costate_equations(t,p,tx,xdata,tu,udata,dHdx)
ut = interp1(tu,udata,t);
xt = interp1(tx,xdata,t);
x1 = xt(1); x2 = xt(2);
dpdt = zeros(2,1);
dpdt = -dHdx(p(1),p(2),ut,xt(1),xt(2)); 
end