num_pts = 200;

x = linspace(-5,5,num_pts);
y = linspace(-5,5,num_pts);
[X,Y] = meshgrid(x,y);
Z = rosen(X,Y);

pts = newton_method([-2,3]);

contour(X,Y,Z, 50)
hold on
plot(pts(1,:),pts(2,:),'-or')

% Actual Newton logic
function pts = newton_method(x)
    %Controlling parameter
    tol = 1e-4;
    maxiter = 500; %Probably far too big
    pts = zeros(2,maxiter); %Preallocate to avoid growth

    pts(:,1) = x(:); %Each column will hold a point
    x = x(:); %Make it be a column vector
    
    d = - hessian(x) \ grad(x);
    eta = norm(d);
    
    iters = 1;
    while eta > tol
        step = armijo(x, d, grad(x)); % Compute step size
        
        x = x + step * d; % Update
        iters = iters + 1;
        pts(:,iters) = x;
        
        if iters > 500
            error('Passed maximal iteration count.')
        end

        % Preparing next iteration
        d = - hessian(x) \ grad(x);
        eta = norm(d);    
    end
    pts = pts(:,1:iters);
end

% For computing step size
function step = armijo(x, d, g)
    % Controlling parameters
    s = 1.0;
    beta = 0.5;
    rho = 0.0001;
    m = 0;
    step = s;
    while (rosen_vec(x) - rosen_vec(x + step * d)) < -rho*s*g'*d
        m = m + 1;
        step = beta * step;
        if m > 200
            error('Unable to find a good step size.')
        end
    end
end

% We use the Rosenbrock function
function f = rosen(x, y)
    f = 100 * (y - x.^2).^2 + (1 - x).^2;
end

% For the iterations we want a version which takes a vector
function f = rosen_vec(x)
    f = rosen(x(1),x(2));
end

function g = grad(x)
    g = [400*x(1)^3 - 400*x(1)*x(2) + 2*x(1) - 2; ...
         -200*x(1)^2 + 200*x(2)];
end

function H = hessian(x)
    H = [1200 * x(1)^2 - 400*x(2) + 2, -400*x(1);
         -400*x(1), 200];
end