clear; clc;
figure(3); clf; set(gcf, 'Color', 'w');
figure(2); clf; set(gcf, 'Color', 'w');
figure(1); clf; set(gcf, 'Color', 'w');

%% Augmented Lagrangian method

% Cost f(x) and single constraint h(x) = 0
f  = @(x) x(1) + x(2);
h  = @(x) x(1)^2 + x(2)^2 - 2;

x0 = [1; 0];
mu = 0;
beta = 1;
niters = 8;
hx = zeros(1, niters);
mus = zeros(1, niters);

figure(1); clf; hold all;
for iter = 1 : niters

    mus(iter) = mu;
    
    [x, info] = alm_subproblem(f, h, beta, mu, x0);

    mu = mu + beta*h(x);

    hx(iter) = h(x);
    x0 = x;
    
    semilogy([info.iter], ...
             [info.gradnorm], '.-', 'LineWidth', 2, 'MarkerSize', 20, ...
             'DisplayName', ['Iter = ', num2str(iter)]);
end

%% Code below is only for plots (no computations)
xlabel('Iteration counter for subproblem');
title('Norm of {\nabla}L_{\beta}(., \mu^k)');
set(gca, 'YScale', 'log');
legend('Location', 'northeast');
set(gca, 'FontSize', 16);

%% More plots as a function of iteration counter k
figure(2); clf;
semilogy(hx, '.-', 'LineWidth', 2, 'MarkerSize', 20);
xlabel('Iteration counter k'); title('Constraint violation h(x_k)');
set(gca, 'FontSize', 16);
figure(3); clf;
plot(mus, '.-', 'LineWidth', 2, 'MarkerSize', 20);
xlabel('Iteration counter k'); title('Parameter \mu^k');
set(gca, 'FontSize', 16);

figure(1); % focus on this one when done