%% Time-Varying Systems and Computations
% Klaus Diepold, Sven Gronauer, Simon Schumann
%
%% Optimal Tracking: row-row-row your boat
% This script explores optimal tracking in boat rowing, emphasizing the use of the pseudo-inverse in an application and the advantages of dynamic programming in this setting.
%% 1. Model of the world
% Remember from the script, our goal is to cross a river with segmentwise constant flux, such that we reach our destination point with least rowing effort. The x coordinate is defined along the flux. The river is divided in n=4 segments for simplicity.
%
% The system dynamics can be described as
%             x_{i+1} = x_i + \delta_i - d_i,
% with x representing the postion, \delta the drift and d the rowing impact.
% Firstly, let's define the drift \delta and the precision M to reach the target. The cost term N_i of rowing is propotional to the drift.

% Natural drift delta
delta_i = [1,2,4,1];
delta_t = sum(delta_i);
% Effort for rowing 
alpha = 1; 
N_i = alpha*delta_i;
N = diag(N_i); 
% Precision reaching destiny
M = 4

% The goal is to reach the target with minimal cost, hence let's construct the cost Cas
%             C_4 = \sum_{i=0}^{3} N_i^2 d_i^2 + M^2 x_4^2.
%  The goal is to find the vector d that minimize the cost.
%
%% 2. The global solution
% As we know how the system evolves over all timesteps as well as the costs associated, we can solve the problem globally via the Moore-Penrose Inverse.
% We formulate the equation Sx=b with S=   \pmatrix{N \cr ME^\prime }, x = \hat{d} and b = -\pmatrix{0 \cr M\delta_t }

E = ones(4, 1); % column vector of 1's
d_hat = zeros(4, 1);
S = [N; -M*E'] % System matrix
b = [zeros(4,1); -M*delta_t] 

% Question: What are the dimension of S, N, x and b?
%%  2.1 With Moore-Penrose Inverse
% For solving this least-squares problem, we calculate the Moore-Penrose inverse of the non-singular system matrix as S^\dagger = \left(N^2 + M^2 EE^\prime\right)^{-1} \left[\matrix{ N & -ME }\right].
% We can now solve for the optimal control \hat{d}:

S_inv = inv(N^2+M^2*(E*E'))*[N, -M*E];
d_hat_glob_inv = S_inv*b

% In addition, let's look to the solutions Matlab procudes with its integrated inv() function.

d_hat_glob_inv_matlab = inv(N^2+M^2*(E*E'))*M^2*E*delta_t

% Did you expect these results?
% How do the results correlate to the velocities of the river?
%
%% 2.2. Sherman Morrison formula
% In this case we can also use the Sherman Morrsion formula (what are the prerequisites for that?) for an explicit matrix inversion and derive the solution as a function of N.

K = delta_t/((1/M^2) + sum(1./N_i.^2))
d_hat_sher = (K./N_i.^2)'

% Looking at the results, we see that the rowing effort N_i^2\hat{d_i}^2 is equal for all segments.
%
%% 2.3 Verification of control inputs
% Finally let's have a look at how the position of the boat elvolves over time. Hence, we use the state space equation to calculate our boat's positions.

x = zeros(5, 1); % Initialize state vector with 5 states, including x_{k+1}
for k = 1:1:4
    x(k+1) = x(k) + delta_i(k) - d_hat_sher(k); % Apply system dynamics equation
end

% Now, we can plot it.

figure;
plot(0:4, x, 'o-'); % Indices from 0 to 4 with markers and lines
title('Boatposition over time M = 4');
xlabel('Timestep');
ylabel('Boat position');
grid on;
xticks(0:4);

% Looking at the final sate, we see that the boat reach closes the desired goal. Also for small flux of the river, the rower moves in the negative direction, while for strong directions it does only hinder stronger drift.
% Finally, let's investigate the impact of the penalization factor M.

plot_state_vectors(delta_t, N_i, delta_i)

% What can you conclude from these diagrams?
% What is the major assumption for tracking globally?
%
%% 3 Dynamic Programming
% The global solution to river navigation lacks adaptivity. By converting global controls into a function of the current state, we gain the efficiency of a local solution that requires only the cost of the next move for optimal decision-making, namely Dynamic Programming.
% In dynamic optimization, we use state x_k to calculate the optimal control \hat{d}_k to progress to x_{k+1} using the state evolution equation and the optimal trajectory cost \hat{C}_{k+1}(x_{k+1}})  that depends solely on x_{k+1}. This recursive process, based on minimal sufficient information at each step, allows for the determination of the optimal local control and cost.
%% 3.1 The local cost model
% One of the major tasks in dynamic programming is to find a suitable local cost model, in particular defining a model to quantifiy on how good it is to be in this state regarding future decisions. Hence we introduce a factor Y to adapt the cost C to a local model  \hat{C}_{k}(x_{k}) for each state.

x(1) = 0; 
x = zeros(5, 1);
d_hat_opt = zeros(4, 1);
costs = zeros(4,1);
Y = zeros(4,1); 

% Next, let's inspect the impact of the factor Y for small and large final cost M

M = 0.5;
Y(4) = M.^2;
Y(3) = N_i(4)*Y(4)/sqrt(N_i(4)^2 + Y(4)^2);
Y(2) = N_i(3)*Y(3)/sqrt(N_i(3)^2 + Y(3)^2);
Y(1) = N_i(2)*Y(2)/sqrt(N_i(2)^2 + Y(2)^2);
disp('M = 0.5');
Y
M = 4;
Y(4) = M.^2;
Y(3) = N_i(4)*Y(4)/sqrt(N_i(4)^2 + Y(4)^2);
Y(2) = N_i(3)*Y(3)/sqrt(N_i(3)^2 + Y(3)^2);
Y(1) = N_i(2)*Y(2)/sqrt(N_i(2)^2 + Y(2)^2);
disp('M = 4');
Y

% What does that mean in context to the final offset in each state?
% As we got a fixed start point, it is suitable to do a forward iteration. We setup the algorithm for calculating the desired control input.

for k = 1:4
    d_hat_opt(k) = get_optimal_control(x, Y, N_i, delta_i, k);
    x(k + 1) = x(k) + delta_i(k) - d_hat_opt(k); % Update state
    costs(k) = get_costs(x, Y, N_i, delta_i, k); % Additionally track costs 
end

% Let's see if it works out.

disp('Optimal control sequence d_hat:');
disp(d_hat_opt);
disp('Cost associated with the optimal control sequence C_hat:');
disp(costs);
disp('Boat position associated with the optimal control sequence x:');
disp(x);

%
%% 3.2 The (lower) QR way
% A different perspective, particularly effective for larger problems, involves working with orthogonalization. In the current case, this approach is especially straightforward.

x = zeros(5, 1);
d_hat_opt_qr = zeros(4, 1);

% After applying a QR factorization, we derive the solution directly as we can minimise the cost of our rowing effort explicitly.

for k = 1:4
    sys_matrix = [Y(k), -Y(k); 0 N_i(k)];
    [Q,L] = ql(sys_matrix);
    delta_t_k = sum(delta_i(k:end));
    d_hat_opt_qr(k) = -(L(2,1)/L(2,2))*(x(k)+delta_t_k);
    x(k + 1) = x(k) + delta_i(k) - d_hat_opt_qr(k); % Update state
end

% At the end, we check again to see if the results are valid.

% Display the results
disp('Optimal control sequence d_hat:');
disp(d_hat_opt_qr);
disp('Cost associated with the optimal control sequence C_hat:');
disp(costs);
disp('Boat position associated with the optimal control sequence x:');
disp(x);

%
%% Appendix

function [Q,L] = ql(A)
    [Q, R] = qr(flipud(fliplr(A)));
    Q = flipud(fliplr(Q));
    L = flipud(fliplr(R));
    if L(end, end) < 0
        L(end, :) = -L(end, :);
        Q(:, end) = -Q(:, end);
    end
end

%

function d_hat_k = get_optimal_control(x, Y, N, delta_i, k)
    delta_t_k = sum(delta_i(k:end));
    coeff = Y(k)^2/(N(k)^2+Y(k)^2);
    d_hat_k = coeff*(x(k)+delta_t_k);
end

function cost_model = get_costs(x, Y, N, delta_i, k)
    delta_t_k = sum(delta_i(k:end));
    coeff = N(k)^2*Y(k)^2/(N(k)^2+Y(k)^2);
    cost_model = coeff*(x(k)+delta_t_k)^2;
end

%

function plot_state_vectors(delta_t, N_i, delta_i)
    num_M_values = 5;
    state_vectors = zeros(num_M_values, 5);

    for M = 1:num_M_values

        K = delta_t / ((1/M^2) + sum(1./N_i.^2));
        d_hat_sher = (K./N_i.^2)';
        x = zeros(1, 5);

        for k = 1:4
            x(k+1) = x(k) + delta_i(k) - d_hat_sher(k);
        end
        state_vectors(M, :) = x;
    end

    % Plotting
    figure;
    hold on;
    for M = 1:num_M_values
        plot(0:4, state_vectors(M, :), 'o-', 'DisplayName', ['M = ' num2str(M)]);
    end
    title('Boat Position Evolution for varying M');
    xlabel('Timestep');
    ylabel('Boat position');
    legend('Location', "southoutside");
    grid on;
    xticks(0:4);
    hold off;
end