% ch13_gradient_descent.m
% Companion code for Chapter 13: Optimization and Data Fitting.
%
% Local minimization by gradient descent with backtracking step control.
% Change the objective, its gradient, and the starting point below.
% The routine accepts any number of variables; the optional plot is for two.
% A small gradient indicates a stationary point, not necessarily a minimum.
% Try different starting points for functions with several local minima.

clear; close all; clc;

%% Problem and settings
% Store the variables in a column vector p. Supply a matching gradient.
f = @(p) p(1)^2 + 4*p(2)^2 + 2*p(1)*p(2);
gradf = @(p) [2*p(1) + 2*p(2); 8*p(2) + 2*p(1)];
p0 = [1.8; -1.6];         % Example minimum: p = [0; 0], f = 0

step_initial = 1.0;        % Trial step at the start of each iteration
tol_grad = 1e-6;           % Absolute tolerance on the gradient norm
max_iter = 1000;           % Maximum number of accepted steps
make_plot = true;          % Optional contour plot for two variables
plot_x = [-0.5, 2.0];      % Adjust these limits for a new example
plot_y = [-2.0, 0.5];

%% Minimize
[pmin, fmin, info] = gradient_descent( ...
    f, gradf, p0, step_initial, tol_grad, max_iter);

%% Results
fprintf('\nGradient descent with backtracking\n');
for j = 1:numel(pmin)
    fprintf('p(%d) = %.12g\n', j, pmin(j));
end
fprintf('Objective value:     %.12g\n', fmin);
fprintf('Gradient norm:       %.6e\n', info.gradient_norm);
fprintf('Accepted steps:      %d\n', info.iterations);
fprintf('Function calls:      %d\n', info.function_calls);
fprintf('Gradient calls:      %d\n', info.gradient_calls);
fprintf('Status:              %s\n', info.status);

%% Optional plot (independent of the numerical routine)
if make_plot && numel(pmin) == 2
    [X, Y] = meshgrid(linspace(plot_x(1), plot_x(2), 150), ...
                     linspace(plot_y(1), plot_y(2), 150));
    Z = arrayfun(@(x,y) f([x; y]), X, Y);
    figure('Color', 'w');
    [~, hc] = contour(X, Y, Z, 25, 'LineColor', [0.55 0.55 0.55]);
    hold on;
    hs = plot(p0(1), p0(2), 'ko', 'MarkerSize', 7);
    hm = plot(pmin(1), pmin(2), 'k+', 'MarkerSize', 10, 'LineWidth', 1.5);
    xlabel('p_1', 'FontSize', 14);
    ylabel('p_2', 'FontSize', 14);
    title('Gradient descent', 'FontSize', 14, 'FontWeight', 'normal');
    legend([hc, hs, hm], {'f(p_1,p_2)', 'Starting point', 'Final estimate'}, ...
        'Location', 'best');
    set(gca, 'FontSize', 14);
    axis equal; xlim(plot_x); ylim(plot_y); grid on; box on;
    % Plot evaluations are excluded from the search function-call count.
end

%% Reusable numerical routine
function [p, fp, info] = gradient_descent(f, gradf, p0, step_initial, tol_grad, max_iter)
% Unconstrained local search for a differentiable scalar objective.
% Backtracking halves the trial step until the Armijo decrease test passes.
% Scale variables sensibly: the gradient tolerance depends on their units.
% Copy this routine and its two helpers to gradient_descent.m for reuse.

    validateattributes(p0, {'double'}, {'real','vector','nonempty','finite'});
    validateattributes(step_initial, {'double'}, {'real','scalar','finite','positive'});
    validateattributes(tol_grad, {'double'}, {'real','scalar','finite','positive'});
    validateattributes(max_iter, {'double'}, ...
        {'real','scalar','finite','integer','nonnegative'});

    p = p0(:);
    fp = objective_value(f, p);
    if ~isfinite(fp)
        error('The objective must be finite at the starting point.');
    end
    g = gradient_value(gradf, p);
    nf = 1; ng = 1; iter = 0;
    armijo = 1e-4;
    max_backtrack = 50;
    status = 'Maximum iteration count reached.';

    while norm(g) > tol_grad && iter < max_iter
        alpha = step_initial;
        accepted = false;
        for trial = 1:max_backtrack
            p_trial = p - alpha*g;
            if isequal(p_trial, p)
                break;
            end
            if all(isfinite(p_trial))
                f_trial = objective_value(f, p_trial);
                nf = nf + 1;
                % Reject nonfinite trial values and try a shorter step.
                if isfinite(f_trial) && ...
                        f_trial <= fp - armijo*alpha*sum(g.^2)
                    accepted = true;
                    break;
                end
            end
            alpha = alpha/2;
        end
        if ~accepted
            status = 'Stopped: backtracking could not find a usable step.';
            break;
        end
        p = p_trial;
        fp = f_trial;
        g = gradient_value(gradf, p);
        ng = ng + 1;
        iter = iter + 1;
    end

    converged = norm(g) <= tol_grad;
    if converged
        status = 'Gradient-norm tolerance reached.';
    end
    info = struct('iterations', iter, 'function_calls', nf, ...
        'gradient_calls', ng, 'gradient_norm', norm(g), ...
        'converged', converged, 'status', status);
end

function value = objective_value(f, p)
    value = f(p);
    if ~isnumeric(value) || ~isscalar(value) || ~isreal(value)
        error('The objective must return a real scalar.');
    end
end

function g = gradient_value(gradf, p)
    g = gradf(p);
    if ~isnumeric(g) || ~isreal(g) || ~isvector(g) || ...
            numel(g) ~= numel(p) || any(~isfinite(g(:)))
        error('The gradient must be a finite real vector matching p.');
    end
    g = g(:);
end