%% Gradient Descent Demo
% Demonstrates gradient descent on the Rosenbrock function
% Teaching Assistant to Professor X

clear; clc; close all;

%% Define the Rosenbrock function
% f(x,y) = (1-x)^2 + 100*(y-x^2)^2
% Minimum at (1,1) with f* = 0

f = @(x) (1-x(1))^2 + 100*(x(2)-x(1)^2)^2;

% Gradient of Rosenbrock
grad_f = @(x) [-2*(1-x(1)) - 400*x(1)*(x(2)-x(1)^2);
               200*(x(2)-x(1)^2)];

%% Algorithm parameters
alpha = 0.001;          % Step size (small due to ill-conditioning)
tol = 1e-6;             % Tolerance for gradient norm
max_iter = 50000;       % Maximum iterations

%% Initialize
x0 = [-1.5; 1.5];       % Starting point
x = x0;

% Storage for convergence history
x_history = zeros(2, max_iter+1);
f_history = zeros(max_iter+1, 1);
grad_norm_history = zeros(max_iter+1, 1);

x_history(:,1) = x;
f_history(1) = f(x);
grad_norm_history(1) = norm(grad_f(x));

%% Gradient Descent iteration
fprintf('Gradient Descent on Rosenbrock Function\n');
fprintf('========================================\n');
fprintf('Initial point: (%.2f, %.2f)\n', x0(1), x0(2));
fprintf('Step size alpha = %.4f\n', alpha);
fprintf('Tolerance = %.1e\n\n', tol);

k = 0;
while k < max_iter
    g = grad_f(x);
    grad_norm = norm(g);

    % Check stopping criterion
    if grad_norm < tol
        fprintf('Converged at iteration %d\n', k);
        break;
    end

    % Gradient descent update
    x = x - alpha * g;
    k = k + 1;

    % Store history
    x_history(:,k+1) = x;
    f_history(k+1) = f(x);
    grad_norm_history(k+1) = grad_norm;

    % Print progress
    if mod(k, 5000) == 0
        fprintf('Iteration %5d: f(x) = %.6e, ||grad|| = %.6e\n', ...
                k, f(x), grad_norm);
    end
end

% Trim history arrays
n_iters = k;
x_history = x_history(:, 1:n_iters+1);
f_history = f_history(1:n_iters+1);
grad_norm_history = grad_norm_history(1:n_iters+1);

%% Print final results
fprintf('\n========================================\n');
fprintf('Final Results:\n');
fprintf('  Iterations: %d\n', n_iters);
fprintf('  Final point: (%.6f, %.6f)\n', x(1), x(2));
fprintf('  Final f(x): %.6e\n', f(x));
fprintf('  Final ||grad||: %.6e\n', norm(grad_f(x)));
fprintf('  Optimal point: (1, 1)\n');
fprintf('  Distance to optimum: %.6e\n', norm(x - [1;1]));

%% Create visualization
figure('Position', [100, 100, 1200, 900], 'Color', 'white');

% Subplot 1: Contour plot with iteration path
subplot(2,2,1);
[X, Y] = meshgrid(linspace(-2, 2, 100), linspace(-1, 3, 100));
Z = (1-X).^2 + 100*(Y-X.^2).^2;
contour(X, Y, log10(Z+1), 30);
hold on;
% Plot every 100th point to avoid clutter
step = max(1, floor(n_iters/500));
plot(x_history(1,1:step:end), x_history(2,1:step:end), 'r.-', 'MarkerSize', 4, 'LineWidth', 1);
plot(x0(1), x0(2), 'go', 'MarkerSize', 10, 'MarkerFaceColor', 'g');
plot(1, 1, 'r*', 'MarkerSize', 15, 'LineWidth', 2);
plot(x(1), x(2), 'bs', 'MarkerSize', 10, 'MarkerFaceColor', 'b');
xlabel('x_1'); ylabel('x_2');
title('Contour Plot with Iteration Path');
legend('Contours', 'Path', 'Start', 'Optimum', 'Final', 'Location', 'best');
colorbar;
grid on;

% Subplot 2: Function value convergence
subplot(2,2,2);
semilogy(0:n_iters, f_history, 'b-', 'LineWidth', 1.5);
xlabel('Iteration k');
ylabel('f(x_k)');
title('Objective Function Convergence');
grid on;

% Subplot 3: Gradient norm convergence
subplot(2,2,3);
semilogy(0:n_iters, grad_norm_history, 'r-', 'LineWidth', 1.5);
xlabel('Iteration k');
ylabel('||\nabla f(x_k)||');
title('Gradient Norm Convergence');
grid on;

% Subplot 4: Distance to optimum
dist_to_opt = sqrt(sum((x_history - [1;1]).^2, 1));
subplot(2,2,4);
semilogy(0:n_iters, dist_to_opt, 'm-', 'LineWidth', 1.5);
xlabel('Iteration k');
ylabel('||x_k - x^*||');
title('Distance to Optimum');
grid on;

% Add super title
sgtitle(sprintf('Gradient Descent: Rosenbrock Function (\\alpha = %.4f, %d iterations)', ...
        alpha, n_iters), 'FontSize', 14, 'FontWeight', 'bold');

%% Save figure
print('gradient_descent_result', '-dpng', '-r150');
fprintf('\nFigure saved as gradient_descent_result.png\n');
