% Douglas-Rachford Splitting for Basis Pursuit
% Solves: minimize ||x||_1 subject to Ax = b
% where A is a 100 x 1000 underdetermined matrix

clear all; close all; clc;

%% Problem Setup
m = 100;   % Number of measurements
n = 1000;  % Signal dimension
k = 20;    % Sparsity level of true signal

fprintf('Setting up basis pursuit problem...\n');
fprintf('Matrix size: %d x %d\n', m, n);
fprintf('Sparsity level: %d\n', k);

% Generate random measurement matrix
rng(123);  % For reproducibility
A = randn(m, n);
A = A / sqrt(m);  % Normalize

% Generate sparse ground truth signal
x_true = zeros(n, 1);
supp = randperm(n, k);  % Random support
x_true(supp) = randn(k, 1);  % Random non-zero values

% Generate measurements
b = A * x_true;

fprintf('Ground truth ||x_true||_1 = %.4f\n', norm(x_true, 1));
fprintf('Ground truth ||x_true||_2 = %.4f\n', norm(x_true, 2));

%% Douglas-Rachford Splitting Algorithm
% Problem: min ||x||_1 s.t. Ax = b
%
% Split into f(x) = ||x||_1 and g(x) = indicator{Ax = b}
%
% Douglas-Rachford iteration:
%   y_{k+1} = prox_g(x_k)
%   x_{k+1} = x_k + prox_f(2*y_{k+1} - x_k) - y_{k+1}

% Algorithm parameters
max_iter = 2000;
lambda = 1.0;  % Step size parameter

% Precompute projection matrix for constraint Ax = b
% Projection: P(z) = z - A'*(A*A')^{-1}*(A*z - b)
AAT = A * A';
AAT_inv = inv(AAT);
proj_matrix = eye(n) - A' * (AAT_inv * A);
proj_offset = A' * (AAT_inv * b);

% Projection operator onto {x: Ax = b}
prox_g = @(z) proj_matrix * z + proj_offset;

% Soft thresholding operator (proximal of ||.||_1)
soft_threshold = @(z, thresh) sign(z) .* max(abs(z) - thresh, 0);
prox_f = @(z) soft_threshold(z, lambda);

% Initialize
x = zeros(n, 1);  % Initial point
y = x;

% Storage for tracking
error_x = zeros(max_iter, 1);        % ||x_k - x_true||_2
error_y = zeros(max_iter, 1);        % ||y_k - x_true||_2
obj_value = zeros(max_iter, 1);      % ||y_k||_1
constraint_residual = zeros(max_iter, 1);  % ||A*y_k - b||_2

fprintf('\nRunning Douglas-Rachford algorithm...\n');
tic;

for iter = 1:max_iter
    % Douglas-Rachford iteration
    y = prox_g(x);                    % Project onto constraint
    z = 2*y - x;                       % Reflected point
    w = prox_f(z);                     % Soft threshold
    x = x + w - y;                     % Update

    % Compute metrics
    error_x(iter) = norm(x - x_true, 2);
    error_y(iter) = norm(y - x_true, 2);
    obj_value(iter) = norm(y, 1);
    constraint_residual(iter) = norm(A*y - b, 2);

    % Display progress
    if mod(iter, 200) == 0 || iter == 1
        fprintf('Iter %4d: ||y-x_true||_2 = %.6f, ||y||_1 = %.4f, ||Ay-b|| = %.2e\n', ...
                iter, error_y(iter), obj_value(iter), constraint_residual(iter));
    end
end

elapsed_time = toc;
fprintf('Completed in %.2f seconds\n', elapsed_time);

%% Results
fprintf('\n=== Final Results ===\n');
fprintf('Ground truth ||x_true||_1 = %.4f\n', norm(x_true, 1));
fprintf('Recovered    ||y||_1      = %.4f\n', obj_value(end));
fprintf('Error ||y - x_true||_2    = %.6f\n', error_y(end));
fprintf('Constraint residual       = %.2e\n', constraint_residual(end));
fprintf('Support recovery: %d/%d correct\n', ...
        sum(abs(y) > 0.1 & abs(x_true) > 0), k);

%% Visualization
figure('Position', [100, 100, 1400, 900]);

% Plot 1: Error vs iterations
subplot(2, 3, 1);
semilogy(1:max_iter, error_y, 'b-', 'LineWidth', 1.5);
grid on;
xlabel('Iteration');
ylabel('Error ||y_k - x_{true}||_2');
title('Reconstruction Error');

% Plot 2: Objective value
subplot(2, 3, 2);
plot(1:max_iter, obj_value, 'r-', 'LineWidth', 1.5);
hold on;
plot([1, max_iter], [norm(x_true, 1), norm(x_true, 1)], 'k--', 'LineWidth', 2);
grid on;
xlabel('Iteration');
ylabel('||y_k||_1');
title('Objective Value (ℓ1 norm)');
legend('Computed', 'True', 'Location', 'best');

% Plot 3: Constraint residual
subplot(2, 3, 3);
semilogy(1:max_iter, constraint_residual, 'g-', 'LineWidth', 1.5);
grid on;
xlabel('Iteration');
ylabel('||Ay_k - b||_2');
title('Constraint Residual');

% Plot 4: True signal
subplot(2, 3, 4);
stem(1:n, x_true, 'k', 'LineWidth', 1.5, 'MarkerSize', 4);
grid on;
xlabel('Index');
ylabel('Amplitude');
title('True Sparse Signal');
xlim([1, n]);

% Plot 5: Recovered signal
subplot(2, 3, 5);
stem(1:n, y, 'b', 'LineWidth', 1.5, 'MarkerSize', 4);
grid on;
xlabel('Index');
ylabel('Amplitude');
title('Recovered Signal (y_k)');
xlim([1, n]);

% Plot 6: Comparison (zoomed on support)
subplot(2, 3, 6);
[~, sorted_idx] = sort(abs(x_true), 'descend');
indices = sorted_idx(1:min(50, n));
stem(indices, x_true(indices), 'ko', 'LineWidth', 1.5, 'MarkerSize', 6);
hold on;
stem(indices, y(indices), 'bx', 'LineWidth', 1.5, 'MarkerSize', 6);
grid on;
xlabel('Index');
ylabel('Amplitude');
title('Signal Comparison (Top 50 components)');
legend('True', 'Recovered', 'Location', 'best');

sgtitle('Douglas-Rachford Splitting for Basis Pursuit', 'FontSize', 14, 'FontWeight', 'bold');

%% Additional convergence plot (first 2000 iterations)
figure('Position', [150, 150, 1000, 600]);

subplot(2, 2, 1);
semilogy(1:max_iter, error_y, 'b-', 'LineWidth', 2);
grid on;
xlabel('Iteration', 'FontSize', 12);
ylabel('||y_k - x_{true}||_2', 'FontSize', 12);
title('Reconstruction Error (ℓ2 norm)', 'FontSize', 13);
set(gca, 'FontSize', 11);

subplot(2, 2, 2);
semilogy(1:max_iter, error_x, 'r-', 'LineWidth', 2);
grid on;
xlabel('Iteration', 'FontSize', 12);
ylabel('||x_k - x_{true}||_2', 'FontSize', 12);
title('Error in Auxiliary Variable x', 'FontSize', 13);
set(gca, 'FontSize', 11);

subplot(2, 2, 3);
plot(1:max_iter, obj_value, 'Color', [0.8, 0.4, 0], 'LineWidth', 2);
hold on;
yline(norm(x_true, 1), 'k--', 'LineWidth', 2, 'Label', 'True ||x||_1');
grid on;
xlabel('Iteration', 'FontSize', 12);
ylabel('||y_k||_1', 'FontSize', 12);
title('Objective Value Evolution', 'FontSize', 13);
set(gca, 'FontSize', 11);

subplot(2, 2, 4);
semilogy(1:max_iter, constraint_residual, 'g-', 'LineWidth', 2);
grid on;
xlabel('Iteration', 'FontSize', 12);
ylabel('||Ay_k - b||_2', 'FontSize', 12);
title('Constraint Satisfaction', 'FontSize', 13);
set(gca, 'FontSize', 11);

sgtitle(sprintf('Douglas-Rachford Convergence Analysis (%d iterations)', max_iter), ...
        'FontSize', 14, 'FontWeight', 'bold');

fprintf('\nPlots generated successfully!\n');
