%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% (CHEM-E7225) ADVANCED PROCESS CONTROL D
%   #02_Extra - Newton's method
%
%   CasADi is a tool to quickly implement and efficiently solve dynamic 
%   optimisation (optimal control) problems. You can also think of it as 
%   a general toolkit for gradient-based optimisation
%
%   CasADI is open-source, LGPL-licensed
%   Download from https://casadi.org  (current version: 3.5.5)
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
close all; clear; clc

addpath('~/Documents/MATLAB/casadi')    % --> addpath('~/change_me_with_path_to/casadi-v3.5.5')
import casadi.*

%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% EXAMPLE: ROOT-FINDING WITH CASADI

% General variables
N = 1;                                  % Input size (Set this to match the example below)
x = MX.sym('x',N,1);                    % R^N

%% Examples of functions to be zero-ed and their Jacobian 
% As symbolic (matrix, MX) expressions

F = exp(x) - 2;                       % R -> R | x^* is log 2
%
% F = (x-3)^3;                          % R -> R | x^* is 3  
%
% F = 1/x - 3;                          % R -> R | x^* is 1/3
%
% F = [ x(1) - x(2) + 1;                % R^2 -> R^2
%       x(2) - x(1)^2 - 1               % x^* are (0,1) and (1,2)
%     ];
%
% F = [ x(1)^2 + x(2)^2 - 5;            % R^2 -> R^2  
%       x(2) - 3*x(1) + 5               % x^* are (2,1) and (1-2)
%     ];
%
% F = [ x(1)^2 + x(2)^2 - 1;            % R^2 -> R^2        
%       sin(pi/2*x(1)) + x(2)^3
%     ];
%
% F = [ exp(x(1)^2 + x(2)^2) - 1;       % R^2 -> R^2
%       exp(x(1)^2 - x(2)^2) - 1 
%     ];
%
% F = [ x(1)^5 + x(2)^3 + x(3)^4 + 1;   % R^3 -> R^3
%       x(1)^2 + x(2) + x(3);
%       x(3)^2 - 1
%     ];

% Function and Jacobian (as CasADI functions)

Ff = Function('Ff', {x}, {F});                % R^N -> R^N
Jf = Function('Jf', {x}, {jacobian(F,x)});    % R^N -> R^(N x N)

%% The exact Newton, with untuned step-size
%  x_k+1 = x_k + J_f(x_k)^(-1) * f(x_k)  
%  From random x0

alpha = 1.0;                            % Step-size
K = 25;                                 % Max number of iterations
R = 5;                                  % Max number of repetitions

X  = nan(N,K,R);                        % Initialise approximations {xk}
Ex = nan(K,R);                          % Initialise `error test` on x      
Ef = nan(K,R);                          % Initialise `error test' on F

for r = 1:R
    x0 = randn(N,1);                % Pick an initial approximation x_0 
    X(:,1,r) = x0;                  % Set initial approximation

    for k = 1:K
        A = Jf(X(:,k,r));           % Compute the Jacobian at xk
        b = -Ff(X(:,k,r));          % Compute function at xk
        deltaX = inv(A)*b;          % Compute the Newton update

        X(:,k+1,r) = X(:,k,r) + alpha * full(deltaX);
        
        Ex(k,r) = log10( norm(X(:,k+1,r) - X(:,k,r)) ); 
        Ef(k,r) = log10( full(norm(Ff(X(:,k+1,r)))) );
    end
    
    X_sol(:,r) = X(:,end,r);
end

display(X_sol)      % Displays the solution obtained in each repetition

%% Plotting the results
figure(); 
ms = 4.0; lw = 1.0; fs = 24;

subplot(2,4,1:4)
 plot(1:K, Ex', '.-', 'MarkerSize', ms)
 ylabel('$\log(\|x_{k+1} - x_{k}\|)$','Interpreter','Latex')
 xlabel('$k$','Interpreter','Latex')
 set(gca,'TickLabelInterpreter','Latex','FontSize',fs)

subplot(2,4,5:8)
 plot(1:K, Ef', '.-', 'MarkerSize', ms)
 ylabel('$\log(\|F(x_{k+1}) - 0\|)$','Interpreter','Latex')
 xlabel('$k$','Interpreter','Latex')
 set(gca,'TickLabelInterpreter','Latex','FontSize',fs)

% -- --