% run wavelet constrained regularization problem with one kind of wavelets at different lambdas with l curve
% Sergey Voronin
close all; clear all;

addpath('utils/');
addpath('wavelets/');
addpath('../solvers_new1/');


% make matrix
fprintf('making matrix..\n');
m = 1200; n = 1500; p = min(m,n);
noise_frac = 0.001;
S = logspace(0,-1.5,p);
S = diag(S);
mat_filename = ['data/mat1.mat']
A = make_normal_random_matrix_specified_decay(m,n,S,mat_filename);
load(mat_filename);


% make vector x... 
fprintf('make x and rhs..\n');
x1 = linspace(-2,2,round(n/2))'; 
x2 = linspace(2,-2,round(n/2))';
x = [x1; x2];
for i=1:n
    if (i < n/4) || (i > 3*n/4)
        if rand < 0.5
            x(i) = x(i) + 0.30*rand;
        else
            x(i) = x(i) - 0.30*rand;
        end
    end
end

% bunch of wide poles - a different x to try (probably better with haar)
%sparsityX = 0.03;
%xref = sprand(n,1,sparsityX);
%x = zeros(n,1);
%mystep_vals = [-50:5:50];
%mystep_vals_ind = 1;
%for i=1:length(x)
%	%if xref(i) ~= 0
%	if mod(i,100) == 0
%		bl = round(5*rand); % block length
%		val = mystep_vals(mystep_vals_ind);
%		mystep_vals_ind = mystep_vals_ind + 1;
%		if mystep_vals_ind < 1 || mystep_vals_ind > length(mystep_vals)
%			mystep_vals_ind = 1;
%		end
%		if ((i + bl) < length(x)) && ((i-bl) > 1)
%			for bl_ind=-bl:bl
%				x(i+bl_ind) = val;
%			end
%		end
%	end
%end
%x = real(x); % filter out any imaginay stuff
%

% make b 
b = A*x;

% add some noise
temp=randn(m,1);%normal distribution
noise = noise_frac*norm(b)/norm(temp)*temp;
b = b + noise;


% fista solution
fprintf('solving..\n');

% set up range of lambdas and parameters
Mtb = haar_1d(A'*b);
lambda_max = max(abs(Mtb))/1.2;
lambda_min = max(abs(Mtb))/5e7;
num_lambdas = 50;
wav_tol = 1e-3; % threshold final wavelet coefficients below this value

lambdas = zeros(num_lambdas,1);
S = (log(lambda_max) - log(lambda_min))/(num_lambdas - 1);
for i=1:num_lambdas
    lambda = exp(log(lambda_max) - S*(i-1));
    lambdas(i) = lambda;
end



lambda_fracs = lambdas/max(abs(Mtb));
num_ticks = num_lambdas;
lambda_fracs_str = {};
for i=1:num_ticks
    if i==1 || i == round(num_ticks/2) 
        lambda_fracs_str{i} = sprintf('%.2e',lambda_fracs(i));
    elseif i==num_ticks
        lambda_fracs_str{i}= sprintf('%.2e',lambda_fracs(i));
    else
        lambda_fracs_str{i} = '';
    end
end


fprintf('run continuation scheme\n');
[x_sol,L_curve_xpts,L_curve_ypts,percent_errors,residual_norms] = continuation_scheme_with_wavelets(A,b,x,lambda_max,lambda_min,num_lambdas,wav_tol);


% compare norm(x) and norm(x_sol), they should be similar
fprintf('final norms and error:\n');
normx = norm(x)
normxsol = norm(x_sol)
percent_error = 100*norm(x_sol - x)/norm(x)

% now you can plot x vs x_sol and etc, as before


% write inputs
%write_mat_to_binary(A,'data/fista1/A.bin');
%write_vec_to_binary(x,'data/fista1/x.bin');
%write_vec_to_binary(b,'data/fista1/b.bin');
fp = fopen('data/fista1/params.txt','w');
fprintf(fp,'lambda_max = %f\n', lambda_max);
fprintf(fp,'lambda_min = %f\n', lambda_min);
fprintf(fp,'num_lambdas = %d\n', num_lambdas);
fprintf(fp,'wav_tol = %f\n', wav_tol);
fclose(fp);


noise_norms = residual_norms;
noise_norms(:) = norm(noise); 


close all;
h = figure(1);
set(h,'PaperPositionMode', 'auto');
%set(gcf,'units','normalized','position',[0 0 1 1]);
hold on;
plot(residual_norms,'r','linewidth',3);
plot(noise_norms,'b--','linewidth',3);
xlim([1,num_lambdas]);
hold off;
legend('residual norm','target level');
tit = title('RESIDUAL NORM vs PARAMETER');
xl = xlabel('lambda fraction');
yl = ylabel('||Ax - b||_2');
set(gca,'XTick',[1:num_ticks]);
set(gca,'XTickLabel',lambda_fracs_str);
set(gca,'fontsize',18);
print('-depsc','residual_norm_curve.eps');


h = figure(2);
set(h,'PaperPositionMode', 'auto');
%set(gcf,'units','normalized','position',[0 0 1 1]);
hold on;
plot(percent_errors,'r','linewidth',3);
xlim([1,num_lambdas]);
ylim([0,100]);
hold off;
tit = title('ERROR vs PARAMETER');
xl = xlabel('lambda fraction');
yl = ylabel('% error');
set(gca,'XTick',[1:num_ticks]);
set(gca,'XTickLabel',lambda_fracs_str);
set(gca,'fontsize',18);
print('-depsc','percent_error_curve.eps');


h = figure(3);
set(h,'PaperPositionMode', 'auto');
%set(gcf,'units','normalized','position',[0 0 1 1]);
plot(L_curve_xpts,L_curve_ypts,'b','linewidth',3);
tit = title('L CURVE');
%plot(fliplr(L_curve_xpts),fliplr(L_curve_ypts));
xl = xlabel('log ||Ax - b||_2');
yl = ylabel('log ||w||_1');
%xlim([2.2,3.5]);
ylim([-5,10]);
%set(gca,'XTick',[1:num_lambdas]);
set(gca,'fontsize',18);
print('-depsc','L_curve.eps');


h= figure(4);
set(h,'PaperPositionMode', 'auto');
%set(gcf,'units','normalized','position',[0 0 1 1]);
hold on;
plot((get_curvature(L_curve_xpts, L_curve_ypts)), 'g','linewidth',3);
xlim([1,num_lambdas]);
hold off;
tit = title('CURVATURE vs PARAMETER');
xl = xlabel('lambda fraction');
yl = ylabel('curvature');
set(gca,'XTick',[1:num_ticks]);
set(gca,'XTickLabel',lambda_fracs_str);
set(gca,'fontsize',18);
print('-depsc','curvature_curve.eps');


h = figure(5);
set(h,'PaperPositionMode', 'auto');
%set(gcf,'units','normalized','position',[0 0 1 1]);
hold on;
plot(x_sol, 'b','linewidth',2);
xlim([1,length(x)]);
ylim([-3,3]);
hold off;
tit = title('RECOVERED X');
xl = xlabel('i');
yl = ylabel('x(i)');
set(gca,'fontsize',18);
print('-depsc','wavelet_transformed_x1.eps');



h = figure(6);
set(h,'PaperPositionMode', 'auto');
%set(gcf,'units','normalized','position',[0 0 1 1]);
hold on;
plot(x, 'b','linewidth',2);
xlim([1,length(x)]);
ylim([-3,3]);
hold off;
tit = title('ORIGINAL X');
xl = xlabel('i');
yl = ylabel('x(i)');
set(gca,'fontsize',18);
print('-depsc','original_x1.eps');


