clear;close all; clc;
tic
%% load the fully sampled myocardial perfusion cartesian dataset
load cart_comp
x=r(:,:,1:70);
IDEAL = r; 
x=double(x);
[n1,n2,n3]=size(x);

m=n1*n2;
n=n3;
x= reshape(x,m,n);
%% load the k-t radial sampling pattern - undersample by ~7.5 fold using radial sampling 
% A frame has 12 radial lines uniformly spaced while other frames
% are rotated by small random angles
load ('mask_12.mat');
mask = fftshift(fftshift(mask,1),2);
S=find(mask~=0);

%% A and At operators are the forward and backward Fourier sampling operators respectively
A = @(z)A_fhp3D(z,S,n1,n2,n3);
At=@(z)At_fhp3D(z,S,n1,n2,n3);


%% first guess
b = A(x);

x_init = At(b);
x_init1=reshape(x_init,n1,n2,n3);
MSE(x_init1,IDEAL)
%% r denotes the number of temporal basis functions in the dictionary
r =45,
%% The algorithm parameters are all specified in opts.
opts.outer =30;% The iterations of the outer loop
opts.inner =40;% The iterations of the inner loop
opts.lambda1=0.0020/(n1*n2);% The regularization parameter (on the l1 norm of U) 
opts.lambda2=1e-8/(n1*n2);% The regularization parameter (on the l2 norm of V)
opts.beta1=1e-2; % continuation parameter for the l1 norm; initialize it



%% Initialize the U and V matrices; (spatial weights/coefficients and temporal bases)
% V - random matrix; U from the CG subproblem based on random V
earray=[];rmse_un=[];earrayout=[];
V = double(rand(r,n)); 
V_init = V;

%  CG(conjugate gradient) algorithm to solve the U initiaization subproblem:

[U,earray_u] = xupdateUal(b,A, At,zeros(m,r,'double'), V,0,opts,1e-10,30);

iii=1;succ_error=[]; X_est=U*V; X_est = reshape(X_est,n1,n2,n3);opts.beta1=1./max(abs(U(:))); 
for out = single(1: opts.outer)
    out
   for in = single(1:opts.inner)
    %   in
       
       %% L subproblem; L - auxillary variable (Shrinkage of spatial weights)
       X_iter=X_est;
       tic
        Z = U;
        L_1 = (abs(Z)-1/opts.beta1); 
        L_1 = L_1.*(L_1>0); 
        L_2 = abs(Z)+(abs(Z)<1e-12);
        L = L_1.*Z./L_2;
          

       %% U subproblem - CG algorithm
  
       [U,earray_u] = xupdateUal(b,A, At,zeros(size(U)), V,L,opts,1e-4,20);

       
       %% V subproblem - CG algorithm
       [V,earray_v] = xupdateVprob(b,A, At,zeros(size(V)), U,opts,1e-4,20);
       
       
       %% COST Calculations
       dc = A(U*V)-b; % data consistency
       reg_U = sum(abs(U(:))); % l1 norm on the spatial weights
       reg_V = sum(abs(V(:)).^2); % l2 norm on the dictionary
       % plot the cost   
       cost = sum(abs(dc(:)).^2) +opts.lambda1*reg_U+opts.lambda2*reg_V;
       earray = [earray,cost];
          X_est = U*V;
        X_est = reshape(X_est,n1,n2,n3);
        
        
        
        
        rmse_un =[rmse_un,MSE(X_est,IDEAL)];
       TIME(iii)=toc;
        iii=iii+1;
       if in>1
        if (abs(earray(end)-earray(end-1))/abs(earray(end))) < 1e-3;
            break; % Break the inner loop if the cost stops decreasing
        end
      
       end
        
        %% plot the results during while the algorithm is converging
       
        
       figure(10); 
       subplot(2,3,1); imagesc(abs(reshape(U(:,3),n1,n2))); title('Example spatial weight')
       %imagesc(abs(reshape(L(:,3),n1,n2))); title('Example auxillary variable')
       subplot(2,3,4); plot(abs(V(3,:))); title('Example temporal basis');
       subplot(2,3,3); imagesc(abs(X_est(:,:,24))); title('A spatial frame: reconstruction');%reshape(L(:,1),n1,n2,1))); title('L');
       subplot(2,3,6); plot(earray); title('Decreasing cost v/s iteration');
       subplot(2,3,5); plot(rmse_un);title('means square error v/s iteration'); hold on; hold off;
       pause(0.01);
   end
   nuc_norm  = X_est; % At low \beta_{1}, the solution is an equivalent min. nuclear norm (low rank) solution
   opts.beta1 = opts.beta1*40; % update the continuation parameter
   earrayout=[earrayout, earray(end)];
   if  out>1
   if (abs(earrayout(end)-earrayout(end-1))/abs(earrayout(end))) < 1e-8  ;
           break; % Break the outer loop if the cost stops decreasing
   end
   end
end



   %% X_est is the final reconstructed image
   close all;
   figure(1);
   subplot(3,2,1); imagesc(abs(IDEAL(:,:,33))); title('IDEAL: a spatial frame');
   subplot(3,2,2); imagesc(abs(squeeze(IDEAL(100,:,:)))); title('IDEAL: image time series');
   subplot(3,2,3); imagesc(abs(X_est(:,:,33))); title('BCS recon: a spatial frame');
   subplot(3,2,4); imagesc(abs(squeeze(X_est(100,:,:)))); title('BCS: image time series');
   subplot(3,2,5); plot(earray); title('Cost');
   subplot(3,2,6); imagesc(abs(mask(:,:,1))); title('Sampling pattern for one frame');
   