%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%% Main function 
%%%%%%
%%%%%% An example for JSENSE method
%%%%%%
%%%%%% Written by: Jinhua Sheng, University of Wisconsin - Milwaukee
%%%%%% Citation: L. Ying and J. Sheng, "Joint image reconstruction and 
%%%%%%           sensitivity estimation in SENSE (JSENSE)", Magn. Reson.
%%%%%%           Med., vol. 57, pp. 1196-1202, June 2007
%%%%%% Created on Oct. 22, 2005
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

clear all;  close all; fclose('all');


ReduceFactor=4;  % Facter: R  = 1 - 8
R=ReduceFactor;

inter_num_VP = 5; 
inter_num = 30 ;
ACSL=16;%32;           %  Number of ACS lines 
order=7;               % Order of polynomial
FOV=256;               % Matrix size
row=FOV;
column=FOV;
alpha=0.0;%0.0015;
JSENSE_tukey_window = cosine_taper_window(128,77,40,2,20,2);

% ****************************************************
%load recon_images;
load GE_human_brain; recon_images=Img;clear Img;
% *****************
a=ones(FOV,1);
b=[a -a];
c=repmat(b,[1 FOV/2]);

%figure;imshow(abs(recon_images(:,:,1)),[]);

% **************************
[D1,D2,CoilNum]=size(recon_images);
for s = 1 : CoilNum
  
    Img(:,:,s)=rot90(recon_images(:,:,s),-1);
end
clear recon_images kspace_data
% ****************************************************
s_poly=zeros(order+1,order+1,CoilNum);
Img0=sqrt(sum([abs(Img)].^2,3));
[mask]=get_mask(Img0);
Img_NMSE=Img0/mean(mean(abs(Img0)));

figure;imshow(abs(rot90(Img0,-1)),[0,0.7*max(max(abs(Img0)))]);title('SoS');

% *******************************
%for s = 1 : CoilNum
%    kspace_data(:,:,s)=fft2(fftshift(Img(:,:,s)));
%    Img_data(:,:,s)=ifft2(kspace_data(:,:,s).*c);
%end 
for s = 1 : CoilNum  
    kspace_data(:,:,s)=fftshift(fft2(Img(:,:,s)));
  %  Img_data(:,:,s)=ifft2(kspace_data(:,:,s).*c);
end  

WeightingFunctions_full=zeros(D2,D2,CoilNum);
for s=1:CoilNum
   WeightingFunctions_full(:,:,s)=(Img(:,:,s)./(Img0+eps));
end
figure;imshow(abs(rot90(WeightingFunctions_full(:,:,1),1)),[]);title('Sensitivity from full data for Channel 1');

% ****************************************
Img_red(:,:,:)=Img(1:ReduceFactor:D2,:,:);

DL=floor(D2/2+0.5)+1-floor(ACSL/2+0.5);
DH=floor(D2/2+0.5)+ACSL-floor(ACSL/2+0.5);
if (DL-1)/R==floor((DL-1)/R)
   DM= floor((DH-1)/R)-floor((DL-1)/R)+1;
else
   DM= floor((DH-1)/R)-floor((DL-1)/R); 
end
if ACSL==0
   DL=floor(D2/2+0.5)+1;
   DH=DL;   DM=0;
end

trajectory=sort(union([1:ReduceFactor:D2],DL:DH));
index=trajectory.';
k0=floor((size(index,1)+1)/2);
area = zeros(size(index,1),1);
area(k0) = 1;
for k = 1 : size(index,1)
    if k > k0
       area(k) = index(k)-index(k-1); 
    end
    if k < k0
       area(k) = index(k+1)-index(k);
    end
end
Density=zeros(D2,D2);
Density(trajectory,:)=area*ones(1,D2);
%figure; imshow(Density,[]);
% ***************************************************


% ***************************************************
% Initial estimated maps
% ***************************************************
center_data(1:DH-DL+1,:,:)=kspace_data(DL:DH,:,:);
initial_map=estimate_maps(ACSL,DL:DH,permute(center_data,[2 1 3]),D2,D2);
for s=1:CoilNum
    WeightingFunctions(:,:,s)=permute(initial_map(:,:,s),[2 1]);
end

figure;imshow(abs(rot90(WeightingFunctions(:,:,1),1)),[]);title('Estimated Sensitivity from ACS for Channel 1');


%Matrix initialization.
kspace_full=kspace_data;
% ******
InitImg=zeros(D2,D2);
nencode=ACSL;
ndim=ceil(D2/ReduceFactor)*ReduceFactor;
acs_line_loc=DL:DH;
KspaceDataWeighted_full=kspace_full;
KspaceDataWeighted=zeros(D2,D2,CoilNum);
KspaceDataWeighted(1:ReduceFactor:D2,:,:)=KspaceDataWeighted_full(1:ReduceFactor:D2,:,:);
[tukey_window,tukey_window_red]=filter_2D(D2,D2,ReduceFactor);
for s = 1 : CoilNum
    kspace_full_GN(:,:,s)=kspace_full(:,:,s).*tukey_window;
end
SY_CG0=ReducedSample(kspace_full,ReduceFactor,D2,ACSL,DL,DH,DM);
SY=ReducedSample(kspace_full_GN,ReduceFactor,D2,ACSL,DL,DH,DM);
%SY_2D=ReducedSample_2D(kspace_full,ReduceFactor,D2,ACSL,DL,DH,DM);
% ************************************************************

% ************************************************************
%   Conjugate Gradient SENSE as the initial value for JSENSE
% ************************************************************
KspaceDataWeighted(acs_line_loc,:,:)=KspaceDataWeighted_full(acs_line_loc,:,:);
Ls = regularization(D2);
ImgRecon_CG=SENSEArbitrary_regul_GN(KspaceDataWeighted,WeightingFunctions,InitImg,trajectory,Density,Ls,alpha,inter_num);
%ImgRecon_CG_NMSE=(ImgRecon_CG*sum(sum(abs(mask)))/sum(sum(abs(ImgRecon_CG.*mask))));
ImgRecon_CG_NMSE=ImgRecon_CG/mean(mean(abs(ImgRecon_CG)));

for s = 1 : CoilNum
    Img_WF(:,:,s)=ImgRecon_CG(:,:).*WeightingFunctions(:,:,s);
    k_space_full_CG(:,:,s)=fftshift(fft2(Img_WF(:,:,s)));
end
SY_CG=ReducedSample(k_space_full_CG,ReduceFactor,D2,ACSL,DL,DH,DM);
error_CG=sum(abs(abs(SY_CG0)-abs(SY_CG)).^2)/sum(abs(SY_CG0).^2)*100;
CG_NMSE=sum(sum(abs(abs(ImgRecon_CG_NMSE)-abs(Img_NMSE)).^2))/sum(sum(abs(Img_NMSE).^2))*100;

figure, imshow(abs(rot90(ImgRecon_CG,-1)),[0,0.7*max(max(abs(ImgRecon_CG)))]); title('CG-SENSE');  

% ************************************************************
%Generate weighting functions and the weighted images.
% Downsampling by a factor ReduceFactor
k_space_red = zeros(floor(D2/ReduceFactor),D2,CoilNum);
k_space_red(:,:,:)=kspace_full(1:ReduceFactor:D2,:,:);
% **************************************************
xx=sum(1:D2)/D2^2; yy=sum(1:D2)/D2^2;
% ************************************************************
%   Sensitivity estimation based on the CG-SENSE reconstruction
% **************************************************************
D1=D2/R;
for s = 1 : CoilNum
    dv = (s-1)*(D1+ACSL-DM)*D2;
    dk_space(:,1)=SY(dv+1:dv+D2*(D1+ACSL-DM),1);
    poly_b=PolyfitIFFT_GN(xx,yy,D1,D2,R,ACSL,DL,DH,DM,order,ImgRecon_CG,dk_space);  
    s_poly(:,:,s) = reshape(poly_b,[order+1,order+1]);        
end
% ****************************************************
% JSENSE: alternating between image reconstruction and sensitivity
% estimation
% ****************************************************
clear kspace_data;
for inter = 1 : inter_num_VP
    inter
    [ImgRecon_JSENSE,s_poly]=JOINT_SENSEArbitrary_regul_JSENSE(KspaceDataWeighted,R,xx,yy,D1,D2,CoilNum,ACSL,DL,DH,DM,SY,mask,s_poly,order,trajectory,Density,Ls,alpha,inter_num,inter);
    ImgRecon_JSENSE_NMSE=ImgRecon_JSENSE/mean(mean(abs(ImgRecon_JSENSE)));
    JSENSE_NMSE=sum(sum(abs(abs(ImgRecon_JSENSE_NMSE)-abs(Img_NMSE)).^2))/sum(sum(abs(Img_NMSE).^2))*100;
    figure, imshow(abs(rot90(ImgRecon_JSENSE,-1)),[0,0.7*max(max(abs(ImgRecon_JSENSE)))]); title(['JSENSE after ',num2str(inter),' iterations' ]);
end



