clear
rng(1,'twister');
%close all

%% Generate a random sym pos def matrix
n = 10;
B = rand(n,n);
%get a random symmetric matrix
B = 0.5*(B + B');
[V,D] = eig(B);
%make its n eigs > 0 by picking from uniform pd [1,n^2]
Dp = randi(n^2,n,1);
%Dp = [1:n];
%Dp = [1,1,1,5,5,5,5,10,10,10];
A = V * diag(Dp) * V';
%prepare a col vector b for use in the quadratic functional
b = rand(n,1);

%% Prepare the quadratic functional and sim params
f = @(x) 0.5 * x' * A * x - b' * x;
gradf = @(x) A * x - b;
%get the true soln
x_true = A\b;
%pick a random starting point 
x0 = rand(n,1);
%sim params
max_iter = 50;
tol = 1e-6; %how low we want the gradient to be

%% Solve using gradient descent and exact step size
%x = x - alpha * gf, where alpha = grad(f)' * grad(f)/(grad(f)'*A*grad(f))
x = x0;
x_SD = [x];
gf = gradf(x);
gf_log = [norm(gf)];
%set initialization params
k = 0;

while norm(gf) > tol && k < max_iter
    alpha = gf'*gf/(gf'*A*gf); %step length
    x_new = x + alpha * (-gf); %move
    gf = gradf(x_new); %gradient for next step
    gf_log = [gf_log, norm(gf)];
    x = x_new;
    k = k+1;
    disp(['SD: k=' num2str(k) ', gf=' num2str(norm(gf))]);
    x_SD = [x_SD, x];
end

figure;
subplot(2,2,1); 
plot(log(gf_log),'b','LineWidth',2); grid;
xlabel('iters');ylabel('Log($\|\nabla(f)\|$)','Interpreter','latex');title('LS convergence');
k_SD = k;


%% Optimize the same function using conjugate gradient method
%x = x + alpha p, where alpha = r'*r/(p'*A*p)
x = x0;
r = gradf(x); 
p = -r;
x_CG = [x];

gf = gradf(x);
gf_log = [norm(gf)];
k = 0;

while norm(gf)>tol && k < max_iter
    alpha = r'*r/(p'*A*p);
    x_new = x + alpha * p;
    %new x found, update params for next iter
    r_new = r + alpha * A * p;
    beta = r_new'*r_new/(r'*r);
    p = -r_new + beta *p;
    r = r_new;
    x = x_new;
    %book keeping
    gf = gradf(x_new);
    gf_log = [gf_log, norm(gf)];
    k = k+1;
    disp(['CG: k=' num2str(k) ', gf=' num2str(norm(gf))]);
    x_CG = [x_CG, x];
end
hold;
plot(log(gf_log),'r','LineWidth',2);
legend(['SD ' num2str(k_SD) ' iters'],['CG ' num2str(k) ' iters']);

subplot(2,2,2);
if n==2
    plot(x0(1),x0(2),'ok','MarkerFaceColor','k')
    hold on;
    plot(x_true(1),x_true(2),'og','MarkerFaceColor','g')
    plot(x_SD(1,:),x_SD(2,:),'b-x')
    plot(x_CG(1,:),x_CG(2,:),'r-x')
    grid; title('Trajectory, SD (blue) v/s CG (red)');
else
   plot(sort(eig(A),'descend'),'x-','LineWidth',2); grid; 
   title('Eigenvalues of A');
end

%% See how the residuals went
subplot(2,2,3);
plot(log(vecnorm(x_SD-repmat(x_true,[1,1+k_SD]))),'b-x','LineWidth',2); grid;
ylabel('$\|x_k-x^*\|$','Interpreter','latex'); xlabel('iters'); title('SD');

subplot(2,2,4);
plot(log(vecnorm(x_CG-repmat(x_true,[1,1+k]))),'r-x','LineWidth',2); grid;
ylabel('$\|x_k-x^*\|$','Interpreter','latex'); xlabel('iters'); title('CG');

