% Load variables t_data and X_data
load('data.mat');

% Initial guess for parameters
k0 = [3.5; 1.5]; 

% Initial state of the system
x0 = [X_data(1,:)'; 0; 0; 0; 0];

% Initialize the path and step size
n_iter = 30;
k_path = zeros(n_iter, 2);
k_path(1,:) = k0;
step_size = 0.1;

%% Optimize k using gradient descent
for i = 2:n_iter
    
    % Current paremeters
    k_i = k_path(i-1, :);
    
    % ODE right-hand side
   odefun = @(t,x) ode_extended(t,x,k_i);

    % Solve the system
   [~, X_out] = ode45(odefun, t_data, x0);
   
   % Solved of dynamic variables
   x_1 = X_out(:,1);
   x_2 = X_out(:,2);
   
   % Solved sensitivities
   s_1 = X_out(:,[3,5]);
   s_2 = X_out(:,[4,6]);
   
   % Compute gradient of the objective using the sensitivities
   dSSE = [2*sum((-X_data(:) + [x_1; x_2]) .* s_1(:));
            2*sum((-X_data(:) + [x_1; x_2]) .* s_2(:))];
   dSSE = dSSE/norm(dSSE);
   
   % Move towards the direction of the gradient
   k_path(i,:) = k_i - step_size*dSSE';

end


%% Plot the optimization path on top of the contour from problem 2
G = 60;
h = linspace(1,6,G);
[K1, K2] = meshgrid(h);
SSE = zeros(size(K1));
for i = 1:G
    for  j = 1:G
        k_ij = [K1(i,j), K2(i,j)];
        SSE(i,j) = compute_SSE(k_ij, t_data, X_data);
    end
end
figure;
contour(K1, K2, log(SSE), 30);
colorbar;
xlabel('k_1');
ylabel('k_2');
title('log(SSE)');
hold on;

% Plot the optimization path
plot(k_path(:,1), k_path(:,2), 'r.-', 'MarkerFaceColor', 'r');