clear all;
clc;

%% Create the system parameters

a = 1.05;
b = 0.01;
q = 1;
r = 1;
x0 = 10;
S = 5;
N = 100;

%% Simulate the system

% Run the function and store the control and performance index
[K,P] = fhopt(a,b,q,r,S,N);

% Create the state variable 
x = zeros(1,N);
x(1) = x0;

% Simulate
for k = 1:N-1
   
    u(k) = -K(k)*x(k);
    x(k+1) = a*x(k) + b*u(k);
    
end

% Plot the state evolution 
figure(1);
plot(x, 'o-', 'LineWidth', 2);
set(gca, 'FontSize',15,'fontname','times','xlim',[1, N],'ylim',[min(x)-0.3,max(x)+0.3]);
grid on;
xlabel('$$\mathbf{k}$$','Interpreter','Latex');
ylabel('$$\mathbf{x_k}$$','Interpreter','Latex');
title('$$\textbf{State Evolution}$$','Interpreter','Latex')

% Plot the input gain K 
figure(2);
plot(K, 'o-', 'LineWidth', 2);
set(gca, 'FontSize',15,'fontname','times','xlim',[1, N],'ylim',[min(K)-0.3,max(K)+0.3]);
grid on;
xlabel('$$\mathbf{k}$$','Interpreter','Latex');
ylabel('$$\mathbf{K_k}$$','Interpreter','Latex');
title('$$\textbf{Input Gain}$$','Interpreter','Latex')

% Plot P  
figure(3);
plot(P, 'o-', 'LineWidth', 2);
set(gca, 'FontSize',15,'fontname','times','xlim',[1, N],'ylim',[min(P)-0.3,max(P)+0.3]);
grid on;
xlabel('$$\mathbf{k}$$','Interpreter','Latex');
ylabel('$$\mathbf{P_k}$$','Interpreter','Latex');

% The state does not converge to zero, because S is finite. This means
% that, for the optimal controller, it yields better performance (i.e.,
% lower cost J) to make less effort to keep the state small by applying
% very small amounts of control close to the end of finite horizon N.
% Although the state is growing close to the end of N=100 steps, control
% effort also decreases and compensates for the penalization of terminal
% state xN*S*xN.

%% Simulate the system with updated terminal cost 

Sup = 500;

% Run the function and store the control and performance index
[K,P] = fhopt(a,b,q,r,Sup,N);

% Create the state variable 
x = zeros(1,N);
x(1) = x0;

% Simulate
for k = 1:N-1
   
    u(k) = -K(k)*x(k);
    x(k+1) = a*x(k) + b*u(k);
    
end

% Plot the state evolution 
figure(4);
plot(x, 'o-', 'LineWidth', 2);
set(gca, 'FontSize',15,'fontname','times','xlim',[1, N],'ylim',[min(x)-0.3,max(x)+0.3]);
grid on;
xlabel('$$\mathbf{k}$$','Interpreter','Latex');
ylabel('$$\mathbf{x_k}$$','Interpreter','Latex');
title('$$\textbf{State Evolution}$$','Interpreter','Latex')

% Plot the input gain K 
figure(5);
plot(K, 'o-', 'LineWidth', 2);
set(gca, 'FontSize',15,'fontname','times','xlim',[1, N],'ylim',[min(K)-0.3,max(K)+0.3]);
grid on;
xlabel('$$\mathbf{k}$$','Interpreter','Latex');
ylabel('$$\mathbf{K_k}$$','Interpreter','Latex');
title('$$\textbf{Input Gain}$$','Interpreter','Latex')

% Plot P  
figure(6);
plot(P, 'o-', 'LineWidth', 2);
set(gca, 'FontSize',15,'fontname','times','xlim',[1, N],'ylim',[min(P)-0.3,max(P)+0.3]);
grid on;
xlabel('$$\mathbf{k}$$','Interpreter','Latex');
ylabel('$$\mathbf{P_k}$$','Interpreter','Latex');

% Now, the state converges to zero, since we penalize the terminal state
% very heavily. The optimal controller, therefore, will try to make the
% state at terminal time k=100 very small such that the terminal cost term
% xN*S*xN is small. When S is very high, this multiplication dominates the
% value of control terms, and therefore the state of the system is driven
% more towards zero.


%% The function for generating and storing the optimal control and 
... performance index

function [K,P] = fhopt(a,b,q,r,S,N)
% Function to compute the optimal control 
    
    K = zeros(N-1,1);
    P = zeros(N,1);
    
    P(N) = S;
    for k = N-1:-1:1
        
        K(k) = (a*b*P(k+1))/(r + P(k+1)*b^2);
        P(k) = q + (r*P(k+1)*a^2)/(r+P(k+1)*b^2);
    end

end