% Newton's Method Demo for Optimization
% Teaching Assistant Demo for Professor X
% Demonstrates quadratic convergence on the Rosenbrock function

clear; clc; close all;

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

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

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

% Hessian
hess_f = @(x) [2 - 400*(x(2) - 3*x(1)^2), -400*x(1);
               -400*x(1),                   200];

%% Newton's Method Implementation
x0 = [0.5; 0.5];           % Initial guess (closer to solution for stability)
tol = 1e-10;               % Tolerance for gradient norm
max_iter = 50;             % Maximum iterations

% Storage for convergence history
x_hist = zeros(2, max_iter+1);
f_hist = zeros(max_iter+1, 1);
grad_norm_hist = zeros(max_iter+1, 1);

% Initialize
x = x0;
x_hist(:,1) = x;
f_hist(1) = f(x);
grad_norm_hist(1) = norm(grad_f(x));

fprintf('Newton''s Method for Rosenbrock Function\n');
fprintf('========================================\n');
fprintf('Initial point: (%.4f, %.4f)\n', x0(1), x0(2));
fprintf('Optimal point: (1, 1) with f* = 0\n\n');
fprintf('%5s %15s %15s %15s\n', 'Iter', 'f(x)', '||grad f||', '||x - x*||');
fprintf('%5s %15s %15s %15s\n', '----', '-------', '----------', '----------');

% Newton iterations
iter = 0;
while iter < max_iter
    g = grad_f(x);

    % Check convergence
    if norm(g) < tol
        fprintf('\nConverged: gradient norm < %.1e\n', tol);
        break;
    end

    % Compute Newton direction: solve H * d = -g
    H = hess_f(x);

    % Check if Hessian is positive definite and regularize if needed
    [~, p] = chol(H);
    if p > 0
        % Hessian not positive definite, use modified Newton
        % Add regularization to make it positive definite
        eigvals = eig(H);
        lambda_min = min(real(eigvals));
        tau = max(abs(lambda_min) + 0.1, 1e-3);
        H = H + tau * eye(2);
        fprintf('  (Hessian regularized at iter %d, tau=%.2e)\n', iter, tau);
    end

    d = -H \ g;  % Newton direction

    % Line search with backtracking (Armijo condition)
    alpha = 1.0;
    c1 = 1e-4;
    rho = 0.5;
    ls_iter = 0;
    max_ls_iter = 50;  % Prevent infinite loop
    while f(x + alpha*d) > f(x) + c1*alpha*(g'*d) && ls_iter < max_ls_iter
        alpha = rho * alpha;
        ls_iter = ls_iter + 1;
        if alpha < 1e-16
            break;
        end
    end

    % Update
    x = x + alpha * d;
    iter = iter + 1;

    % Store history
    x_hist(:, iter+1) = x;
    f_hist(iter+1) = f(x);
    grad_norm_hist(iter+1) = norm(grad_f(x));

    % Print progress
    x_star = [1; 1];
    fprintf('%5d %15.6e %15.6e %15.6e\n', iter, f(x), norm(grad_f(x)), norm(x - x_star));
end

% Trim history arrays
x_hist = x_hist(:, 1:iter+1);
f_hist = f_hist(1:iter+1);
grad_norm_hist = grad_norm_hist(1:iter+1);

fprintf('\n========================================\n');
fprintf('Final solution: (%.10f, %.10f)\n', x(1), x(2));
fprintf('Final f(x): %.6e\n', f(x));
fprintf('Final ||grad f||: %.6e\n', norm(grad_f(x)));
fprintf('Total iterations: %d\n', iter);

%% Create visualization figure
figure('Position', [100, 100, 1200, 400]);

% Subplot 1: Contour plot with iteration path
subplot(1,3,1);
[X, Y] = meshgrid(linspace(-0.5, 1.5, 100), linspace(-0.5, 1.5, 100));
Z = (1 - X).^2 + 100*(Y - X.^2).^2;
contour(X, Y, log10(Z + 1), 30);
hold on;
plot(x_hist(1,:), x_hist(2,:), 'ro-', 'LineWidth', 2, 'MarkerSize', 6, 'MarkerFaceColor', 'r');
plot(x_hist(1,1), x_hist(2,1), 'gs', 'MarkerSize', 12, 'MarkerFaceColor', 'g', 'LineWidth', 2);
plot(1, 1, 'mp', 'MarkerSize', 15, 'MarkerFaceColor', 'm', 'LineWidth', 2);
xlabel('x_1', 'FontSize', 12);
ylabel('x_2', 'FontSize', 12);
title('Newton''s Method Path', 'FontSize', 14);
legend('log_{10}(f+1) contours', 'Newton path', 'Start', 'Optimum', 'Location', 'best');
colorbar;
grid on;
axis equal;

% Subplot 2: Function value convergence
subplot(1,3,2);
semilogy(0:iter, f_hist, 'b-o', 'LineWidth', 2, 'MarkerSize', 6);
xlabel('Iteration', 'FontSize', 12);
ylabel('f(x_k)', 'FontSize', 12);
title('Function Value Convergence', 'FontSize', 14);
grid on;

% Subplot 3: Gradient norm convergence
subplot(1,3,3);
semilogy(0:iter, grad_norm_hist, 'r-s', 'LineWidth', 2, 'MarkerSize', 6);
xlabel('Iteration', 'FontSize', 12);
ylabel('||\nabla f(x_k)||', 'FontSize', 12);
title('Gradient Norm Convergence', 'FontSize', 14);
grid on;

% Add overall title
sgtitle('Newton''s Method on Rosenbrock Function', 'FontSize', 16, 'FontWeight', 'bold');

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

%% Demonstrate quadratic convergence rate
fprintf('\n========================================\n');
fprintf('Quadratic Convergence Analysis:\n');
fprintf('========================================\n');
x_star = [1; 1];
errors = zeros(iter+1, 1);
for k = 1:iter+1
    errors(k) = norm(x_hist(:,k) - x_star);
end

fprintf('%5s %15s %15s\n', 'Iter', '||x_k - x*||', 'Ratio (quadratic)');
fprintf('%5s %15s %15s\n', '----', '------------', '-----------------');
for k = 1:iter
    if errors(k) > 1e-15
        ratio = errors(k+1) / errors(k)^2;
        fprintf('%5d %15.6e %15.6f\n', k-1, errors(k), ratio);
    end
end
fprintf('\nQuadratic convergence: ||e_{k+1}|| / ||e_k||^2 approaches constant\n');
